* OutOfMemoryError}.
   *
   * @param iterator the iterator to return a partitioned view of
   * @param size the desired size of each partition (the last may be smaller)
   * @return an iterator of immutable lists containing the elements of {@code iterator} divided into
   *     partitions
   * @throws IllegalArgumentException if {@code size} is nonpositive
   */

  }

  // In theory we can use PartitionedCall if the XLA cluster does not have any
  // stateful operations.  However, for now we choose to be conservative since
  // we don't have any evidence that choosing a stateless partitioned call helps
  // for performance.
  ops::StatefulPartitionedCall call(
      root.WithOpName("stateful_partitioned_call"), args, n->output_types(),

// CHECK-SAME:   topology = "topology"
// CHECK:        tf_device.return %[[CLUSTER]]#0, %[[CLUSTER]]#1
// CHECK:      return %[[REPLICATE]]#0, %[[REPLICATE]]#3


// Test replication with model parallelism using partitioned resource inputs.
// The cluster will be wrapped in a `tf_device.cluster` first and then by a
// replicate.
// TPUPartitionedInputV2 nodes would be inside the replicate but outside the
// cluster.

message RollingUpdateStatefulSetStrategy {
  // Partition indicates the ordinal at which the StatefulSet should be partitioned
  // for updates. During a rolling update, all pods from ordinal Replicas-1 to
  // Partition are updated. All pods from ordinal Partition-1 to 0 remain untouched.
  // This is helpful in being able to do a canary based deployment. The default value is 0.
  optional int32 partition = 1;

      This transformation pass optimizes the op execution order of the ops in
      the model.
  }];
  let constructor = "CreateOptimizeOpOrderPass()";
}

def PartitionedTopologicalSortPass : Pass<"tfl-partitioned-topological-sort", "mlir::func::FuncOp"> {
  let summary = "Re-sort execution order such that delegated ops stay together";
  let constructor = "CreatePartitionedTopologicalSortPass()";
    let description = [{

	"":               "RollingUpdateStatefulSetStrategy is used to communicate parameter for RollingUpdateStatefulSetStrategyType.",
	"partition":      "Partition indicates the ordinal at which the StatefulSet should be partitioned for updates. During a rolling update, all pods from ordinal Replicas-1 to Partition are updated. All pods from ordinal Partition-1 to 0 remain untouched. This is helpful in being able to do a canary based deployment. The default value is 0.",

message RollingUpdateStatefulSetStrategy {
  // Partition indicates the ordinal at which the StatefulSet should be partitioned
  // for updates. During a rolling update, all pods from ordinal Replicas-1 to
  // Partition are updated. All pods from ordinal Partition-1 to 0 remain untouched.
  // This is helpful in being able to do a canary based deployment. The default value is 0.
  optional int32 partition = 1;

                        "number of logical devices : {2}",
                        input_index, tiled_input_size, num_cores_per_replica));
    };

    // If input is already partitioned using the `tf.TPUPartitionedInputV2` op,
    // only replicated sharding is supported where i-th operand to
    // `tf.TPUPartitionedInputV2` op is input to the i-th logical device.
    if (auto partitioned_input =

type RollingUpdateStatefulSetStrategy struct {
	// Partition indicates the ordinal at which the StatefulSet should be partitioned
	// for updates. During a rolling update, all pods from ordinal Replicas-1 to
	// Partition are updated. All pods from ordinal Partition-1 to 0 remain untouched.
	// This is helpful in being able to do a canary based deployment. The default value is 0.
	Partition *int32 `json:"partition,omitempty" protobuf:"varint,1,opt,name=partition"`

	"":               "RollingUpdateStatefulSetStrategy is used to communicate parameter for RollingUpdateStatefulSetStrategyType.",
	"partition":      "Partition indicates the ordinal at which the StatefulSet should be partitioned for updates. During a rolling update, all pods from ordinal Replicas-1 to Partition are updated. All pods from ordinal Partition-1 to 0 remain untouched. This is helpful in being able to do a canary based deployment. The default value is 0.",