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

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