options.device_ordinal = 0;
  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);

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

                                 "/job:localhost/replica:0/task:0");
    device_ = device.get();
    SetDevice(device_type, std::move(device));

    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_fns{
        UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
    device_context_ = core::RefCountPtr<DeviceContext>(
        new PjRtDeviceContext(shape_fns, /*use_pjrt_tensor_buffer=*/true));

    // Get the host allocator.

  XlaHelpers::ShapeRepresentationFn shape_representation_fn =
      IdentityShapeRepresentationFn();
  XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns{
      UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
  return PopulateResultIOInfo(module_op, arg_shapes, /*use_tuple_args=*/false,
                              /*use_resource_updates_for_aliases=*/false,

    // have one program running on each core at the same time.
    options.use_global_compute_stream = true;
    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);