std::move(compiler_client));
}

std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
GetShapeDeterminationFns() {
  XlaHelpers::ShapeRepresentationFn shape_representation_fn =
      [](const TensorShape&, DataType, bool, XlaLayoutPreference) {
        return xla::Shape();
      };
  XlaShapeLayoutHelpers::LayoutPreferenceFn 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)
          << ",allow_cpu_custom_calls=" << options.allow_cpu_custom_calls
          << ",populate_resource_manager=" << options.populate_resource_manager

    auto layout_preference = shape_determination_fns.layout_preference_fn(
        arg_shapes[i].shape, arg_dtype, std::nullopt);
    TF_ASSIGN_OR_RETURN(xla_shape,
                        shape_determination_fns.shape_representation_fn(
                            arg_shapes[i].shape, arg_dtype,
                            /*use_fast_memory=*/false, layout_preference));

    // Rewrite layout with sharding, if sharding is set.

          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(),
                            /*fast_mem=*/false, layout_preference));

    // The device tensor should always be fresh.

      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(),
                          /*fast_mem=*/false, layout_preference));
  const xla::Layout* device_layout = &(shape.layout());

  // from failures.
  Status EnsureDeviceContextOk() TF_LOCKS_EXCLUDED(mu_);

  // Two convenient methods to get the underlying device context.
  // Get the default device context, created by the first
  // shape_representation_fn.
  absl::StatusOr<DeviceContext*> GetDeviceContextDefault();
  // Get the device context given the index.
  absl::StatusOr<DeviceContext*> GetDeviceContextWithIndex(int index);

  xla::XlaComputation* xla_computation = compilation_result->computation.get();
  *xla_computation = xla::XlaComputation(hlo_proto.hlo_module());

  XlaHelpers::ShapeRepresentationFn shape_representation_fn =
      IdentityShapeRepresentationFn();
  XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns{
      UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};

        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(
                            input->shape(), input->dtype(),
                            /*use_fast_memory=*/false, layout_preference));
    TF_RETURN_IF_ERROR(xla_output->AllocateShapedBuffer(

  options.compilation_device_name = DEVICE_CPU_XLA_JIT;
  options.use_multiple_streams = false;
  XlaShapeLayoutHelpers::ShapeDeterminationFns shape_representation_fns{
      UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
  options.shape_determination_fns = {shape_representation_fns};
  auto device = std::make_unique<XlaDevice>(session_options, options);

  // Setting AcceleratorDeviceInfo because eager runtime relies on the device

    options.use_multiple_streams = true;
    options.allowed_devices = gpu_ids;
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_representation_fns{
        UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
    options.shape_determination_fns = {shape_representation_fns};
    auto device = std::make_unique<XlaDevice>(session_options, options);

    Status status = device->UseAcceleratorDeviceInfo();