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

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

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

    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,

    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(

    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,
    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,

// MLIR
//   native kernels for legalization to HLO.
// return_tuple: when this is true, always create a tuple result for the
//   entry computation.
// shape_determination_fns: Contains layout preference fn and shape
//   representation fn. The two functions are used to determine argument and
//   result shapes.
// custom_legalization_passes: passes to run before the default TF legalization

    bool lower_to_xla_hlo, const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, llvm::StringRef device_type,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,
    bool use_tuple_args, XlaCompiler::CompilationResult* compilation_result,
    std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
    const std::vector<TensorShape>& arg_shapes,