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>
        shape_determination_fns_;

    Metadata(const Metadata&) = delete;

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

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

  Status ThenExecute(Device* device, stream_executor::Stream* stream,

  CHECK(xla_tensor);

  XlaLayoutPreference layout_preference =
      shape_determination_fns_.layout_preference_fn(
          device_tensor->shape(), device_tensor->dtype(), std::nullopt);
  Status status = [&]() -> Status {
    TF_ASSIGN_OR_RETURN(xla::Shape shape,
                        shape_determination_fns_.shape_representation_fn(
                            device_tensor->shape(), device_tensor->dtype(),

    const DeviceType& jit_device_type() const;
    const XlaShapeLayoutHelpers::ShapeDeterminationFns&
    default_shape_determination_fns() const {
      return shape_determination_fns_.at(0);
    }
    const PaddedShapeFn& padded_shape_fn() const { return padded_shape_fn_; }

    bool UseMultipleStreams() const { return use_multiple_streams_; }

   private:
    const int device_ordinal_;
    const DeviceType device_type_;

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

    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(
                          cpu_tensor->shape(), cpu_tensor->dtype(),

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