}

    auto failure_or_converged = InferModuleShape(
        getOperation(), max_iterations_, /*ops_to_skip=*/{}, input_shapes_);
    if (failed(failure_or_converged)) return signalPassFailure();
    if (!failure_or_converged.value()) {
      getOperation().emitError()
          << "shape inference pass did not reach convergence after "
          << max_iterations_;
      return signalPassFailure();
    }
  }

 private:

                                          int64_t max_iterations);

  // Shape propagation for call/control flow ops.
  // Returns a failure() on error, otherwise returns true to indicate that it
  // reached convergence, false otherwise.
  FailureOr<bool> PropagateShapeIntoAttachedFunctions(Operation* op,
                                                      int64_t max_iterations);

  // Shape propagation for region based control flow.

    are only kept in the context rather than the ops folded/IR mutated.
  }];

  let constructor = "TF::CreateTFShapeInferencePass()";

  let options = [
    Option<"max_iterations_", "max-iterations", "int64_t", /*default=*/"10",
           "Maximum shape inference iterations">,
    Option<"input_arg_shapes_", "input-arg-shapes", "std::string", /*default=*/"",

    The Pmap executes body function in parallel for all ranges up to $max_iterations.

    The pseudo code:
      for(int i = 0; i < $max_iterations; i++) {
        body_fn(MlrtFture($tensor_list_or_flow_in[i]),
                MlrtPromise($tensor_list_or_flow_in[i+1]),
                i, i, $invariant_args);
      }

      return $tensor_list_or_flow_in[$max_iterations]
  }];

  let arguments = (ins
    TF_Tensor:$max_iterations,

    The Pmap executes body function in parallel for all ranges up to $max_iterations.

    The pseudo code:
      for(int i = 0; i < $max_iterations; i++) {
        body_fn(MlrtFture($tensor_list_or_flow_in[i]),
                MlrtPromise($tensor_list_or_flow_in[i+1]),
                i, i, $invariant_args);
      }

      return $tensor_list_or_flow_in[$max_iterations]
  }];

  let arguments = (ins

// reached convergence, false otherwise.
// If input shapes are provided, first refines the `main` function using
// InferShapeForFunction.
FailureOr<bool> InferModuleShape(ModuleOp module, int64_t max_iterations = 10,
                                 ArrayRef<TypeID> ops_to_skip = {},
                                 ArrayRef<ArrayRef<int64_t>> input_shapes = {});

  // Run clustering passes until the convergence. Limit the number of iterations
  // to guard from the infinite loop in presence of bugs.
  constexpr int max_iterations = 100;
  for (unsigned i = 0; i < max_iterations; ++i)
    if (!RunClusteringPass(state, policies)) break;

  // Form clusters found by the union-find algorithm.
  llvm::DenseMap<unsigned, Cluster> root_clusters;

  return %2 : tensor<3xf32>
}

// -----

// Test a while to map_fn conversion in which max_iterations are passed
// into the predicate function.

// CHECK-LABEL: @"map/while_cond"

	om.setsIndexer.Add(set)
	ssc.enqueueStatefulSet(set)
	fakeWorker(ssc)
	selector, err := metav1.LabelSelectorAsSelector(set.Spec.Selector)
	if err != nil {
		return err
	}
	iterations := 0
	for (maxIterations < 0 || iterations < maxIterations) && set.Status.ReadyReplicas < *set.Spec.Replicas {
		iterations++

  // Look if there is a function that tries until it converge.
  // Currently unit-test doesn't do multiple tries, so we need this.
  const int max_iterations = 15;
  for (int i = 0; i < max_iterations; ++i) {
    if (failed(applyPartialConversion(func, target, frozenPatterns))) {
      return false;
    }
  }
  return true;
}

void LegalizeTFPass::runOnOperation() {