// [g1, g2, g4, g5, g3]
  auto data = op.input(0);
  auto partitions = op.input(1);
  int32_t num_partitions;
  TF_RETURN_IF_ERROR(
      GetNodeAttr(op.node()->attrs(), "num_partitions", &num_partitions));

  // Note: the shape of the partitions is a prefix of the data shape.
  // shape(partitions) = [5]
  auto partitions_shape = Shape(scope, partitions);
  // We now create a partitions-shaped tensor with integers from

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

    // CHECK-SAME:   module = "ML\EFR{{.*}}"

    %0 = "tf.XlaCallModule"(%arg0, %arg1) {Sout = [#tf_type.shape<?>], dim_args_spec = [], _entry_function = @_stablehlo_main_0, _stablehlo_module_attrs = { mhlo.num_partitions = 1 }, module = "", platforms = [], version = 5 : i64} : (tensor<10xi32>, tensor<10xi32>) -> tensor<10xi32>
    // CHECK: return %[[RESULT]]
    func.return %0 : tensor<10xi32>
  }

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

03\0F%\1F/!!)#\1F\19)\1F\13\15\1D\15G\11\1F\15\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00return_v1\00custom_call_v1\00call_v1\00xla_call_module_serialization.mlir\00arg_attrs\00function_type\00res_attrs\00sym_name\00sym_visibility\00mhlo.num_partitions\00api_version\00backend_config\00call_target_name\00called_computations\00has_side_effect\00operand_layouts\00output_operand_aliases\00result_layouts\00tf.backend_config\00callee\00\00_stablehlo_f\00jax.arg_info\00x\00mhlo.sharding\00{replica...

\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]
```