std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
        shape_determination_fns;

    // If padded_shape_fn is empty, a default implementation that returns
    // the logical on-device shape without padding is used.
    PaddedShapeFn padded_shape_fn;

    // Set of devices to use. This controls which of the devices on the given
    // platform will have resources allocated. For GPUs this will be

        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)),
      use_multiple_streams_(use_multiple_streams) {}

    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns{
        UseNoPreferenceLayoutFn(), &TpuShapeRepresentation};
    options.shape_determination_fns = {shape_determination_fns};
    options.padded_shape_fn = &TpuPaddedShapeFn;
    auto device = std::make_unique<XlaDevice>(session_options, options);

    // The AcceleratorDeviceInfo actually provides information not only for GPU

    DeviceType compilation_device_type) {
  return std::make_unique<XlaDevice::Metadata>(
      /*device_ordinal=*/0, /*platform=*/nullptr, compilation_device_type,
      GetShapeDeterminationFns(), XlaDevice::PaddedShapeFn(),
      /*use_multiple_streams=*/false);
}

std::unique_ptr<PjRtBaseDevice::Metadata> CreatePjRtDeviceMetadata(
    DeviceType compilation_device_type) {

      auto padded_shape_i = padded_shape[1 + i];
      auto block_shape_ints_i = block_shape_ints[i];

      // Compute the outer_shape constant values to infer the reshape.
      if (padded_shape_i == -1 || block_shape_ints_i == -1) {
        outer_shape_ints.push_back(-1);
      } else {
        outer_shape_ints.push_back(padded_shape_i / block_shape_ints_i);
      }
    }