tf_executor.fetch
  }
  func.return
}

// -----

// TODO(b/188046643): Fuse op for partitioned variable within the island.

// Check that we bring in TPUPartitionedInput operand producers.
// CHECK-LABEL: func @fuse_in_partitioned_input_op

	"":               "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.",

  // resource variables and stateful functions.
  std::vector<Operation*> topological_order;
  for (Operation& op : parent_func.getOps())
    if (ops.contains(&op)) topological_order.push_back(&op);

  // Create the partitioned call
  builder.restoreInsertionPoint(insertion_point);
  auto caller = MakeFuncCaller(builder, module.getLoc(), new_func, inputs);

  Block* block = new_func.addEntryBlock();

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.

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

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.

            insert(call.getArgOperands()[entry.getFirst()], source, func_block);
      }
    });
  }
  return result;
}

// Contains cached information for decomposed callee functions for (stateful)
// partitioned call ops.
struct PartitionedCallTensorArrayOpsInfo {
  bool signature_change;
  func::FuncOp decomposed_callee;
  llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<string, 4>>, 4>
      arg_grads;

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.

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.

  func.return %0 : tensor<*x!tf_type.resource<tensor<4xf32>>>
}

// -----

// Tests that the pass lifts resources on two partitioned call ops sharing the
// same callee. The lifting should clone the callee then modify the clone.

// CHECK-LABEL: @cluster_with_partitioned_call
func.func @cluster_with_partitioned_call() -> tensor<f32> {