pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass());
  }

  // Legalize while early to allow further constant folding.
  // TODO(jpienaar): This may not actually matter as we do canonicalization
  // after the legalize below, for now it needs to be below the above passes
  // that work on TF dialect and before inliner so that the function calls in

          for (auto arg : block.getArguments()) {
            ConvertAndWrapUsesInUnrealizedCast(arg, converter, rewriter);
          }
        }
      }
    });
  });
  return success();
}

// Legalize StableHLO portion of program to VHLO, leaves TFL untouched
LogicalResult ApplyStablehloToVhloPatterns(ModuleOp module,
                                           bool is_func_legal) {
  MLIRContext *context = module.getContext();

    func.return
  }
})";

// MLIR which should not legalize at all
static constexpr char kBadMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  func.func @main() -> () {
    %0 = tf.Unknown() -> ()
    func.return %0
  }
})";

// MLIR which should be filtered by the MLIR bridge but fully legalize with the
// combined bridge.

constexpr char kTfLiteInputIndices[] = "_tflite_input_indices";

// Legalize operations in functions.
class LegalizeTFPass : public impl::LegalizeTFPassBase<LegalizeTFPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeTFPass)
  LegalizeTFPass() = default;
  explicit LegalizeTFPass(bool run_tfl_runtime_verification,
                          bool preserve_assert_op) {

  [TF_UnsortedSegmentMaxOp, TFL_UnsortedSegmentMaxOp],
  [TF_UnsortedSegmentMinOp, TFL_UnsortedSegmentMinOp],
  [TF_UnsortedSegmentProdOp, TFL_UnsortedSegmentProdOp]] in {
  def Legalize#UnsortedSegmentOp[0] :
     Pat<(UnsortedSegmentOp[0] $data, $segment_ids, $num_segments),
              (UnsortedSegmentOp[1] $data, (CreateTFCastToInt32Op $segment_ids),

// (either defined in TensorArrayV3 or implied in the first write).
std::unique_ptr<OperationPass<ModuleOp>>
CreateTensorArrayOpsDecompositionPass();

// Create a pass that legalize TFG to TF dialect.
std::unique_ptr<Pass> CreateLegalizeTFGToTFEPass();

// Matches sequence of ops to TensorFlow fused kernels. This pass should not be

                        lower_to_xla_hlo);

  if (lower_to_xla_hlo) {
    // This pass operates on MHLO control flow ops so it should be legalized
    // after the control flow ops are legalized.
    pm.addPass(mlir::mhlo::CreateLegalizeTFCommunicationPass());

    // Everything should be MHLO after this.
    if (!allow_partial_conversion) {
      pm.addNestedPass<mlir::func::FuncOp>(

    "/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count",
    "Counts the failure of legalization of ops", "op_name", "legality");

class LegalizeTF : public impl::LegalizeTFBase<LegalizeTF> {
 public:
  explicit LegalizeTF(bool legalize_chlo,
                      std::optional<StringRef> tf2xla_fallback_device_type,
                      bool prefer_tf2xla) {
    legalize_chlo_ = legalize_chlo;

//   . Legalizes control flow operations.
//   . Decomposes compound resource operations so that the only remaining
//     operations on resource variables are resource reads/writes..
//   . Replaces resource reads/writes with function inputs/outputs and
//     eliminates the use of resource variables.
//   . Legalizes the operations to XLA HLO operations.

        verifyTaskIds()
        verifyTaskDetails(RegisterTaskBuildOperationType, withPath(':', ':foo')).children.empty
        def realize = verifyTaskDetails(RealizeTaskBuildOperationType, withPath(':', ':foo'))
        realize.children.empty
        realize.details.eager == false
    }

    def "op during realize are child ops"() {
        given:
        register('foo')
        register('bar')
        buildFile << """