}

    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=*/"",

// CHECK-LABEL: func.func private @"map/while_cond"
func.func private @"map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
  %0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
  return %0 : tensor<i1>
}

// `max_iterations` are reached. Returns failure if resource ops decomposition
// does not converge after `max_iterations`.
// TODO(prakalps): This can be useful to a lot of other passes in bridge.
// Extract out as a separate utility.
LogicalResult ApplyPatternsLocallyUntilConverged(
    Operation* op_with_regions, FrozenRewritePatternSet& patterns,
    int max_iterations) {
  bool changed = true;

    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 = {});

    int minThresholdPosition = 0;
    // The leftmost position at which the greatest of the k lower elements
    // -- the new value of threshold -- might be found.

    int iterations = 0;
    int maxIterations = IntMath.log2(right - left, RoundingMode.CEILING) * 3;
    while (left < right) {
      int pivotIndex = (left + right + 1) >>> 1;

      int pivotNewIndex = partition(left, right, pivotIndex);

    int minThresholdPosition = 0;
    // The leftmost position at which the greatest of the k lower elements
    // -- the new value of threshold -- might be found.

    int iterations = 0;
    int maxIterations = IntMath.log2(right - left, RoundingMode.CEILING) * 3;
    while (left < right) {
      int pivotIndex = (left + right + 1) >>> 1;

      int pivotNewIndex = partition(left, right, pivotIndex);