(*name_to_node_map)[node->name()] = node;
  }
}

// Strips off the tensor part of the tensor_name to get the node_name.
const string GetNodeNameFromTensorName(string tensor_name) {
  if (tensor_name[0] == '^') {
    tensor_name.erase(0, 1);
  }
  std::vector<string> tensor_name_parts = str_util::Split(tensor_name, ':');
  return tensor_name_parts[0];
}

    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
  );
}

  }

  // Runs and compares the outputs of `tensor_name` on both the
  // `unfrozen_session` and the `frozen_graph_def.
  void RunAndCompareFrozenAndUnfrozenGraphs(Session* unfrozen_session,
                                            const GraphDef& frozen_graph_def,
                                            const string& tensor_name) {
    std::vector<Tensor> unfrozen_outputs;

  } 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) {

  TFE_OpAddInput(op, in, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TF_DeleteStatus(status);

  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(in));
  TFE_OpSetAttrString(op, "tensor_name", "dummy", 5);
  TFE_OpSetAttrString(op, "send_device", send_device.c_str(),
                      send_device.size());
  TFE_OpSetAttrString(op, "recv_device", recv_device.c_str(),

    });
    if (!status.ok()) {
      done(status);
    }
  }
}

void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                             absl::string_view tensor_name,
                                             Device* device, Tensor* cpu_tensor,
                                             StatusCallback done) {
  if (device_tensor->NumElements() == 0) {

  // CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute([[recv]])
  // CHECK: tf_mlrt.await [[result_future]]
  %0 = "tf.Recv"() {__op_key = 0: i32, device = "/device:CPU:0", tensor_name = "tensor", send_device = "/device:CPU:0", send_device_incarnation = 0, recv_device = "/device:CPU:0"} : () -> tensor<f32>