// 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).

  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();