const XlaShapeLayoutHelpers::ShapeDeterminationFns&
    default_shape_determination_fns() const {
      return shape_determination_fns_.at(0);
    }

    const XlaShapeLayoutHelpers::ShapeDeterminationFns&
    shape_determination_fns_at(int i) const {
      return shape_determination_fns_[i];
    }

   private:
    const DeviceType jit_device_type_;
    std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>

          << ",graph_def_version=" << options.graph_def_version
          << ",options.shape_determination_fns.layout_preference_fn?="
          << (options.shape_determination_fns.layout_preference_fn != nullptr)
          << ",options.shape_determination_fns.shape_representation_fn?="
          << (options.shape_determination_fns.shape_representation_fn !=
              nullptr)

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

          use_tuple_args, true, shape_determination_fns, compilation_result,
          custom_legalization_passes, metadata.module_name(),
          lower_to_xla_hlo));

  // Compute how arguments are shared across different cores.
  auto sharding_result =
      tpu::GetShardingInfo(metadata, arg_shapes, shape_determination_fns,
                           arg_core_mapping, per_core_arg_shapes);

    return device_to_device_streams_.at(index).get();
  }
  xla::TransferManager* transfer_manager() const { return transfer_manager_; }
  const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns()
      const {
    return shape_determination_fns_;
  }

  // Returns a device-to-device stream, in round-robin fashion.
  se::Stream* GetDeviceToDeviceStream();

  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),

    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,
    std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    const std::vector<tensorflow::TensorShape>& arg_shapes,
    std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,

// devices using PjRt.
class PjRtDeviceContext : public DeviceContext {
 public:
  explicit PjRtDeviceContext(
      XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
      bool use_pjrt_tensor_buffer = false)
      : shape_determination_fns_(std::move(shape_determination_fns)),
        use_pjrt_tensor_buffer_(use_pjrt_tensor_buffer) {}

  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,

    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(