%1 = "tf.I"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
    func.return %1 : tensor<?xi32>
  }
}

// -----

// Test `tf_device.cluster_func` on TPU with pre-split replicate sharded
// input/output using `tf.TPUPartitionedInputV2` and `tf.TPUPartitionedOutputV2`.