VLOG(2) << "Compiled PJRT executable " << executable->name()
          << " num_replicas " << executable->num_replicas()
          << " num_partitions " << executable->num_partitions();

  return std::move(executable);
}

absl::StatusOr<std::string> PjRtDeviceCompilerClient::SerializeExecutable(
    const xla::PjRtLoadedExecutable& executable) {

    %0:2 = "tf.XlaCallModule"(%arg0, %arg1, %arg2) {
      Sout = [#tf_type.shape<10>], dim_args_spec = [],
      _entry_function = @main_0,
      _stablehlo_module_attrs = { mhlo.num_partitions = 1 }, module = "",
      platforms = [], version = 5 : i64
    } : (tensor<10xf32>, tensor<10xf32>, tensor<6xi32>) -> (tensor<10xf32>, tensor<6xi32>)

\05\0B\03\05\1D\05\06!\03\05\05\01\01\07\04\01\03\03\06\03\01\05\01\00f\051\0F\0B\03!\1B\1D[;\05\1F\15\1D\15\1D%)9\13\15\19\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00multiply_v1\00return_v1\00sym_name\00jax.uses_shape_polymorphism\00mhlo.num_partitions\00mhlo.num_replicas\00jit_jax_model\00arg_attrs\00function_type\00res_attrs\00sym_visibility\00x\00jit(jax_model)/jit(main)/mul\00experimental/users/ypang/lite/convert_ulm.py\00mhlo.sharding\00{replicated}\00jax.result_info\00\00main\00public\00",...

  // Restore the deserialized stablehlo module's attributes to the reconstructed
  // stablehlo module. The stablehlo module's attributes can contain important
  // information such as SPMD num_replicas and num_partitions.
  auto original_stablehlo_module_attrs =
      op->getAttrOfType<DictionaryAttr>(kStablehloModuleAttrsAttrName);
  if (original_stablehlo_module_attrs) {

  std::vector<std::unique_ptr<xla::PjRtBuffer>> execute_outputs;
  std::optional<xla::PjRtFuture<>> future;
  if (executable->num_replicas() != 1 || executable->num_partitions() != 1) {
    TF_ASSIGN_OR_RETURN(
        execute_outputs,
        executable->ExecuteSharded(
            executable_args, device,
            GetPjRtExecuteOptions(device_type,

LogicalResult TPUPartitionedInputV2Op::verify() {
  TPUPartitionedInputV2Op op = *this;

  int num_partitions = 1;
  const mlir::ArrayAttr partition_dims = op.getPartitionDims();
  for (const mlir::Attribute &dim : partition_dims) {
    num_partitions *= dim.cast<IntegerAttr>().getInt();
  }

  const bool is_packed = op.getIsPacked();
  const bool replicated = partition_dims.empty();

For example:

```python
    # Scalar partitions.
    partitions = 1
    num_partitions = 2
    data = [10, 20]
    outputs[0] = []  # Empty with shape [0, 2]
    outputs[1] = [[10, 20]]

    # Vector partitions.
    partitions = [0, 0, 1, 1, 0]
    num_partitions = 2
    data = [10, 20, 30, 40, 50]
    outputs[0] = [10, 20, 50]
    outputs[1] = [30, 40]
```

    xla::ExecuteOptions exe_options;
    exe_options.arguments_are_tupled = false;
    exe_options.untuple_result = true;

    // TODO(b/257548614): currently PJRT is compiled as portable (num_replica =
    // 1 and num_partition = 1). Support multiple partitions case.
    return executable->ExecutePortable(executable_args, pjrt_device,
                                       exe_options);
  }