%1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  // CHECK:      [[COMPUTATION:%.+]] = "tf_device.cluster_func"([[INPUT]])
  %2 = "tf_device.cluster_func"(%1) {func = @computation, use_spmd_for_xla_partitioning = true} : (tensor<i32>) -> tensor<i32>
  // CHECK:      [[OUTPUT:%.+]]:2 = "tf.TPUPartitionedOutputV2"([[COMPUTATION]])
  // CHECK-SAME: _XlaSharding = ""

  let summary = "Rewrites a `tf_device.cluster_func` on TPUs into TPU runtime operations.";

  let description = [{
    This pass rewrites a `tf_device.cluster_func` operation into a sequence of `tf._TPUCompileMlir`
    and `tf.TPUExecute` operations. `tf._TPUCompileMlir` contains a MLIR module that is
    functionally equivalent to the function referenced by `tf_device.cluster_func`.

// RUN: tf-opt %s -tf-device-mark-input-output-aliases | FileCheck %s

// The following tests check if the aliasing pairs are conservatively marked
// correctly. In the following tests tf_device.cluster_func has inputs
// coming from ReadVariableOp and outputs written to a resource using
// AssignVariableOp. If a pair of input-output (say input at index `a` and
// output at index `b`) read and write to the same resource, then those
// input-output pairs alias each other.

    %0 = "tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_resources} : (tensor<!tf_type.resource>, tensor<i32>) -> tensor<i32>
    // CHECK: "tf.XlaLaunch"(%arg1, %arg0) <{function = @func_with_resources, operandSegmentSizes = array<i32: 0, 1, 1>}> : (tensor<i32>, tensor<!tf_type.resource>) -> tensor<i32>

// Tests `tf_device.cluster_func` with missing `step_marker_location` attribute.

module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
  func.func @bad_num_cores_per_replica() {
    // expected-error@+1 {{requires attribute 'step_marker_location'}}

      BuildSpaceToDepth(cluster_func, input, block_size, input_shape);
  cluster_func.setOperand(index, space_to_depth);
  return space_to_depth;
}

// Performs transformation for replicated inputs. Returns true if this is a
// supported case (thus transform happened).
bool HandleHostReplicatedInputs(int64_t index,
                                tf_device::ClusterFuncOp cluster_func,

  module->walk([&](mlir::tf_device::ClusterFuncOp cluster_func) {
    cluster_func_ops.push_back(cluster_func);
  });
  EXPECT_EQ(cluster_func_ops.size(), 1);

  EXPECT_TRUE(mlir::succeeded(tensorflow::EraseClusterFuncs(cluster_func_ops)));

  llvm::SmallVector<mlir::tf_device::ClusterFuncOp, 4> new_cluster_func_ops;
  module->walk([&](mlir::tf_device::ClusterFuncOp cluster_func) {
    new_cluster_func_ops.push_back(cluster_func);
  });

  }
  return v;
}

void AnnotateParameterReplicationPass::runOnOperation() {
  ModuleOp m = getOperation();
  OpBuilder builder(m.getContext());
  m.walk([&](tf_device::ClusterFuncOp cluster_func) {
    auto replicate = cluster_func->getParentOfType<tf_device::ReplicateOp>();
    if (!replicate) return;
    auto mirrored_variable_indices_attr =
        replicate->getAttrOfType<ArrayAttr>(kMirroredVariableIndicesAttr);

    mlir::OpBuilder* builder,
    llvm::SmallVectorImpl<llvm::SmallVector<mlir::Value, 4>>* input_list);

// Extracts a list of OpSharding that represent output sharding configuration of
// `tf_device.cluster`.
mlir::LogicalResult ParseAndValidateOutputSharding(
    int num_cores_per_replica, mlir::tf_device::ClusterFuncOp cluster_func,

func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>