hlo_sharding, /*use_fast_memory=*/false, shape_determination_fns, shape));
  return absl::OkStatus();
}

// Converts arg_shapes to xla::Shape's and store into xla_input_shapes.
Status GetXlaInputShapes(
    mlir::ModuleOp module, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
    bool use_tuple_args,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    std::vector<xla::Shape>* xla_input_shapes) {

  TF_ASSERT_OK_AND_ASSIGN(
      auto shape, options.shape_determination_fns.shape_representation_fn(
                      TensorShape(), DT_FLOAT, false,
                      tensorflow::XlaLayoutPreference::kTpuPreferLinearLayout));
  EXPECT_EQ(shape, xla::Shape());
  EXPECT_EQ(options.shape_determination_fns.layout_preference_fn(
                TensorShape(), DT_FLOAT, std::nullopt),

    mlir::ModuleOp module_op, llvm::StringRef device_type,
    xla::XlaComputation* xla_computation, bool use_tuple_args,
    bool enable_op_fallback, bool return_tuple,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns =
        {},
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes = {},
    llvm::StringRef module_name = llvm::StringRef());

    const auto& shape_determination_fns =
        dst_xla_context->shape_determination_fns();
    XlaLayoutPreference layout_preference =
        shape_determination_fns.layout_preference_fn(
            input->shape(), input->dtype(), std::nullopt);
    TF_ASSIGN_OR_RETURN(xla::Shape shape,
                        shape_determination_fns.shape_representation_fn(

    xla::PjRtClient* pjrt_client,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns
        shape_determination_fns) {
  XlaLayoutPreference layout_preference =
      shape_determination_fns.layout_preference_fn(
          cpu_tensor->shape(), cpu_tensor->dtype(), std::nullopt);
  TF_ASSIGN_OR_RETURN(xla::Shape shape,
                      shape_determination_fns.shape_representation_fn(

    std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
        shape_determination_fns,
    PaddedShapeFn padded_shape_fn, bool use_multiple_streams)
    : device_ordinal_(device_ordinal),
      device_type_(device_type),
      platform_(platform),
      shape_determination_fns_(std::move(shape_determination_fns)),
      padded_shape_fn_(std::move(padded_shape_fn)),

      IdentityShapeRepresentationFn();
  XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns{
      UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
  return PopulateResultIOInfo(module_op, arg_shapes, /*use_tuple_args=*/false,
                              /*use_resource_updates_for_aliases=*/false,
                              shape_determination_fns, compilation_result);
}

      device_to_device_streams_(std::move(device_to_device_streams)),
      client_(client),
      transfer_manager_(client->backend().transfer_manager()),
      shape_determination_fns_(std::move(shape_determination_fns)),
      thread_pool_(thread_pool) {
  CHECK(host_to_device_stream_ != nullptr);
  CHECK(stream_ != nullptr);
}

void XlaDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor,

  class Metadata {
   public:
    Metadata(int device_ordinal, se::Platform* platform,
             const DeviceType& device_type,
             std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
                 shape_determination_fns,
             PaddedShapeFn padded_shape_fn, bool use_multiple_streams);

    // The index of the device on this host.
    int device_ordinal() const;

    se::Platform* platform() const;

    }

    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns;

    absl::Status compilation_status =
        tensorflow::tf2xla::v1::CompileTensorflowGraphToHlo(
            computation, metadata_proto, use_tuple_args,
            shape_determination_fns, arg_shapes, &arg_core_mapping,
            &per_core_arg_shapes, client, &compilation_result);