}

// Creates a 1D string array constant for "tensor_names" input of `RestoreV2`
// op. The `ConstOp` will be created at `builder`'s current insertion point.
TF::ConstOp CreateTensorNamesConst(const ArrayRef<std::string> tensor_names,
                                   OpBuilder& builder) {
  const auto loc = NameLoc::get(builder.getStringAttr("tensor_names"));
  return Create1DStringConst(tensor_names, loc, builder);
}

    const std::vector<std::string>& tensor_names) {
  absl::flat_hash_map<std::string, int> name_to_index;
  for (const auto& index_and_tensor : llvm::enumerate(subgraph.tensors)) {
    name_to_index[index_and_tensor.value()->name] = index_and_tensor.index();
  }

  std::vector<int> indices;
  indices.reserve(tensor_names.size());

  for (const auto& name : tensor_names) {
    auto found = name_to_index.find(name);

  let arguments = (ins
    Arg<TF_StrTensor, [{Must have a single element.  The prefix of a V2 checkpoint.}]>:$prefix,
    Arg<TF_StrTensor, [{shape {N}.  The names of the tensors to be restored. This is tensor name in the model.}]>:$tensor_names,
    Arg<TF_StrTensor, [{shape {N}.  The slice specs of the tensors to be restored.
Empty strings indicate that they are non-partitioned tensors.}]>:$shape_and_slices,
    // VarHandles for restored tensors.

    because the tensors are loaded asynchronously.

  }];

  let arguments = (ins
    TFTensorType:$prefix,
    TFTensorType:$tensor_names,
    TFTensorType:$shape_and_slices,
    Variadic<TFTensorType>:$var_handles,
    TypeArrayAttr: $restored_dtypes,
    DenseBoolArrayAttr: $truncate_in_cast
  );
}

  }];
}

def TF_SaveOp : TF_Op<"Save", []> {
  let summary = "Saves the input tensors to disk.";

  let description = [{
The size of `tensor_names` must match the number of tensors in `data`. `data[i]`
is written to `filename` with name `tensor_names[i]`.

See also `SaveSlices`.
  }];

  let arguments = (ins

  }
}

AssetFileDef CreateAssetFileDef(const StringRef filename,
                                const StringRef tensor_name) {
  AssetFileDef asset_file_def{};
  asset_file_def.set_filename(MaybeStripAssetDirectoryPrefix(filename).str());

  tensorflow::TensorInfo tensor_info{};
  tensor_info.set_name(tensor_name.str());
  *asset_file_def.mutable_tensor_info() = tensor_info;

  return asset_file_def;
}

#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/tf2xla/type_util.h"

namespace tensorflow {

void XlaHostRecvDeviceContext::CopyDeviceTensorToCPU(
    const Tensor* device_tensor, StringPiece tensor_name, Device* device,
    Tensor* cpu_tensor, StatusCallback done) {
  DataType dtype = EncodePrimitiveTypeAsDataType(shape_.element_type()).value();
  TensorShape tensor_shape;

                             bool sync_dst_compute) const override;

  void CopyDeviceTensorToCPU(const Tensor* device_tensor,
                             StringPiece tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override {
    done(errors::Internal("host->device copy not implemented."));
  }

  }

  // Copies `device_memory_base_` with `shape_` into `cpu_tensor`.
  // `device_tensor` is unused.
  void CopyDeviceTensorToCPU(const Tensor* device_tensor,
                             StringPiece tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override;

  void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
                              Tensor* output_tensor,

  } else {
    return first_try_buffer.status();
  }
}
}  // namespace

void PjRtDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                              absl::string_view tensor_name,
                                              Device* device,
                                              Tensor* cpu_tensor,
                                              StatusCallback done) {