// CHECK: [[RI:%.*]] = "tf.TPUReplicatedInput"([[PI_0]], [[PI_1]])
  %pi_0 = "tf.TPUPartitionedInput"(%arg0, %arg1) {_XlaSharding = "", partition_dim = -1 : i64} : (tensor<!tf_type.resource<tensor<10x3xf32>>>, tensor<!tf_type.resource<tensor<10x3xf32>>>) -> tensor<!tf_type.resource<tensor<10x3xf32>>>

// Check that we bring in TPUPartitionedInput operand producers.
// CHECK-LABEL: func @fuse_in_partitioned_input_op
func.func @fuse_in_partitioned_input_op(%arg0: tensor<2x4xf32>, %arg1: tensor<2x4xf32>) {
  tf_executor.graph {
// CHECK: tf_executor.island
// CHECK-NEXT: = "tf.TPUPartitionedInput"
// CHECK-NEXT: = "tf.Const"
// CHECK-NEXT: = "tf.AddV2"

mlir::Operation* NullUnlessSharded(PartitionedOp op) {
  return op.get_XlaSharding() ? op : nullptr;
}

// Returns a TPUPartitionedInput op connected to a `tf_device.cluster_func`
// operand value if it has an XLA sharding. If value is a resource type then
// TPUPartitionedInput op will be connected to a ReadVariable op that feeds into
// a `tf_device.cluster_func`.
mlir::Operation* GetXlaShardingFromOperand(Value value) {

    // expected-error @+1 {{'tf.TPUPartitionedInput' op TF2XLA TPU bridge input check: number of inputs inconsistent. num_cores_per_replica=2 no. of inputs=3}}
    %pi, %c0 = tf_executor.island wraps "tf.TPUPartitionedInput"(%arg0, %arg1, %arg1) {index = 1 : i64} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>

      std::vector<IslandOp> child_ops;
      if (incoming) {
        // Looks at captured operands of `candidate_wrapped_op` to bring special
        // TPU ops such as tf.TPUReplicatedInput and tf.TPUPartitionedInput into
        // the island as well. These ops are brought in only if they do not
        // already have a cluster assigned to them (via `_replication_info`
        // attribute value).

inputs. It reorders `tf.TPUPartitionedInput` (model parallelism) and
`tf.TPUReplicatedInput` (data parallelism) ops. It transforms a DAG where
multiple `tf.TPUPartitionedInput` ops are feeding into a single
`tf.TPUReplicatedInput` into a DAG where multiple `tf.TPUReplicatedInput` ops
are feeding into a single `tf.TPUPartitionedInput`. Transforming the IR in such

    inputs. It reorders `tf.TPUPartitionedInput` (model parallelism) and
    `tf.TPUReplicatedInput` (data parallelism) ops. It transforms a DAG where
    multiple `tf.TPUPartitionedInput` ops are feeding into a single
    `tf.TPUReplicatedInput` into a DAG where multiple `tf.TPUReplicatedInput` ops
    are feeding into a single `tf.TPUPartitionedInput`. Transforming the IR in such

  let description = [{
    Bubbles up sharding configuration from `cluster_func` regions into
    the attributes of `cluster_func`. This is done by parsing the
    `XlaSharding` / `TPUPartitionedOutput` / `TPUPartitionedInput` ops inside
    `cluster_func`.

    For example, given the following `cluster_func` wrapping `func`:

    ```mlir
      func @test(%arg0: tensor<*xi32>) {
        "tf_device.cluster_func"(%arg0) {

// attributes back to legacy attributes.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateConvertToLegacyCompileAndReplicateAttributesPass();

// Creates a pass that converts all TPUPartitionedInput to TPUPartitionedInputV2
std::unique_ptr<OperationPass<func::FuncOp>>
CreateTPUPartitionedOpConversionPass();

std::unique_ptr<OperationPass<ModuleOp>> CreateTPUValidateInputsPass();

    TF_Int64Tensor:$device_ordinal
  );
}

// We must manually define TPUPartitionedInput, TPUPartitionedInputV2,
// TPUPartitionedOutput, and TPUPartitionedOutputV2 since they have an
// optional attribute, _XlaSharding, unlike the TensorFlow definition.
def TF_TPUPartitionedInputOp : TF_Op<"TPUPartitionedInput", [Pure]> {
  let summary = [{
An op that groups a list of partitioned inputs together. This op