public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReduceTypePrecisionPass)
  void runOnOperation() override;
};

void ReduceTypePrecisionPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  patterns.add<CheckRangeAndConvertI8ToI4, SanitizeGatherOpOutputToI4>(
      &getContext());
  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
}

      TryPruneResultDefiningOp(while_op, result);
    }
  });
  // Now eliminate passthrough operands/results with existing canonicalization
  // pattern.
  MLIRContext* context = &getContext();
  RewritePatternSet patterns(context);
  TF::WhileRegionOp::getCanonicalizationPatterns(patterns, context);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    signalPassFailure();
  }
}

struct BatchMatMulToEinsumPass
    : public impl::BatchMatMulToEinsumPassBase<BatchMatMulToEinsumPass> {
  void runOnOperation() override;
};

void BatchMatMulToEinsumPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();

  patterns.add<ConvertTFBatchMatMulToEinsumOp<TF::BatchMatMulOp>,
               ConvertTFBatchMatMulToEinsumOp<TF::BatchMatMulV2Op>>(
      &getContext());

    // ops will be erased during the optimization.
    rewriter.eraseOp(assign_op);
  }
};

void RemoveVariableInitializationByConstPass::runOnOperation() {
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<RemoveVariableAssignmentByConst>(&ctx);

  ModuleOp module_op = getOperation();
  func::FuncOp init_func_op = GetInitializerFunction(

    }
  }
};

#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/cast_bf16_ops_to_f32.inc"

void CastBf16OpsToF32Pass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  RewritePatternSet patterns(ctx);
  auto module_op = getOperation();

  patterns.add<CastBf16OpsToF32>(ctx);
  populateWithGenerated(patterns);

  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {

  rewriter.replaceOpWithNewOp<StorageCastOp>(qbarrier, qbarrier.getType(),
                                             newConstOp);
  return success();
}

void ConvertConstPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();
  auto *context = &getContext();
  patterns.add<QuantizedConstRewrite>(context);
  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
}

    target.addDynamicallyLegalOp<mlir::stablehlo::CustomCallOp>(
        [&](mlir::stablehlo::CustomCallOp op) {
          return op.getCallTargetName() != "stablehlo.composite";
        });

    RewritePatternSet patterns(context);
    patterns.add<ReplaceCustomCallWithComposite>(context);

    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns)))) {

    return success();
  }
};

void MergeFusionWithDequantizePass::runOnOperation() {
  ModuleOp module_op = getOperation();
  MLIRContext* ctx = module_op.getContext();
  RewritePatternSet patterns(ctx);
  patterns.add<MergeFusionWithUniformDequantizePattern>(ctx);
  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {
    signalPassFailure();
  }
}

}  // namespace

    if (add_op) {
      rewriter.replaceOp(add_op, op->getResult(0));
    }

    return success();
  }
};

void GpuOpFusionPass::runOnOperation() {
  func::FuncOp func = getOperation();
  RewritePatternSet patterns(&getContext());
  patterns.add<ReluToFusedBatchNorm>(&getContext());
  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
}

}  // namespace

  }
};

static PassRegistration<ConvertCustomAggregationOpToQuantStatsPass> pass;

void ConvertCustomAggregationOpToQuantStatsPass::runOnOperation() {
  MLIRContext *ctx = &getContext();
  RewritePatternSet patterns(ctx);
  func::FuncOp func = getOperation();

  patterns.add<ConvertCustomAggregationOpToQuantStats>(ctx);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError()