const string& name, std::unique_ptr<tensorflow::Tensor>* out_tensor,
    TF_Status* out_status) const {
  Status status;
  if (reader_ != nullptr) {
    status = reader_->GetTensor(name, out_tensor);
  } else {
    tensorflow::DataType dtype;
    tensorflow::TensorShape shape;
    status = v2_reader_->LookupDtypeAndShape(name, &dtype, &shape);
    if (status.ok()) {
      out_tensor->reset(new Tensor(dtype, shape));

  const TensorSliceReader::VarToDataTypeMap& GetVariableToDataTypeMap() const;

  // Attempts to look up the tensor named "name" and stores the found result in
  // "out_tensor".
  void GetTensor(const string& name,
                 std::unique_ptr<tensorflow::Tensor>* out_tensor,
                 TF_Status* out_status) const;

 private:
  // Uses "v2_reader_" to build "var name -> shape" and "var name -> data type"

  Tensor* CreateHostTensor(const TensorShape& shape,
                           const gtl::ArraySlice<T> data) {
    Tensor* host_tensor =
        new Tensor(host_allocator_, DataTypeToEnum<T>::v(), shape);
    test::FillValues<T>(host_tensor, data);
    tensors_.push_back(host_tensor);
    return host_tensor;
  }

  // Creates a Tensor on device using the device_allocator_
  template <typename T>

  Tensor* CreateHostTensor(const TensorShape& shape,
                           const gtl::ArraySlice<T> data) {
    Tensor* host_tensor =
        new Tensor(host_allocator_, DataTypeToEnum<T>::v(), shape);
    test::FillValues<T>(host_tensor, data);
    tensors_.push_back(host_tensor);
    return host_tensor;
  }

  // Creates a Tensor on device using the device_allocator_
  template <typename T>

        args.device_context = new XlaHostRecvDeviceContext(
            stream, device_memory_base, shape, done_event);

        Tensor host_tensor;
        TF_RETURN_IF_ERROR(
            ctx->rendezvous()->Send(parsed_key, args, host_tensor, false));

        return std::move(done_event);
      };
}

// Provide RecvDeviceMemoryFunction for XLA host callbacks.  This callback

    DefaultValuedOptionalAttr<BoolAttr, "false">:$enable_large_batch_splitting
  );

  let results = (outs
    Res<Variadic<TF_Tensor>, [{The output tensors.}]>:$out_tensors
  );

  TF_DerivedOperandTypeListAttr Tcaptured = TF_DerivedOperandTypeListAttr<1>;
  TF_DerivedOperandTypeListAttr Tin = TF_DerivedOperandTypeListAttr<0>;