}

    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:

    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,

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