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

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

    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(

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

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

    }

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