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