enable_per_channel_quantized_weight_ = enable_per_channel_quantized_weight;
  }

 private:
  void runOnOperation() override;
};

void QuantizePass::runOnOperation() {
  ModuleOp module_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<StableHloQuantization, StableHloQuantizationReverse>(&ctx);

  PopulateCommonQuantizationPatterns(ctx, patterns,

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

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateReduceTypePrecisionPass() {
  return std::make_unique<ReduceTypePrecisionPass>();
}

}  // namespace TFL

  for (const unsigned operand_index : erase_indices.set_bits()) {
    operand_index < num_in_original ? num_in_tensors-- : num_captured_tensors--;
  }

  getOperation()->eraseOperands(erase_indices);
  getOperation()->setAttr(
      operand_segment_size_attr,
      DenseI32ArrayAttr::get(
          getContext(), /*content=*/{num_in_tensors, num_captured_tensors}));
}

void TPURewritePass::runOnOperation() {
  TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(getOperation(), &devices)))
    return signalPassFailure();

  // Collect compilation results.
  llvm::DenseMap<Attribute, SmallVector<TF::TPUCompilationResultOp, 1>>
      compilation_results;
  auto result_init = getOperation().walk([&](TF::TPUCompilationResultOp op) {

 public:
  explicit AnalyzeVariablesPass() {}
  void runOnOperation() override;
};

void AnalyzeVariablesPass::runOnOperation() {
  auto* context = &getContext();
  auto module = getOperation();
  bool legalize_to_tfl = true;

  module.walk([&](Operation* op) {
    // Skip ops that are supported natively by TFLite.
    if (IsSupportedTFLiteResourceOp(op)) return WalkResult::advance();

  // Find all the executes first, since we will mutate the nodes around each
  // execute in the same tf_device.replicate op.
  llvm::SmallVector<tf_device::LaunchOp, 8> execute_launches;
  getOperation().walk([&](tf_device::LaunchOp op) {
    if (op.WrapsSingleOp() &&
        llvm::isa<TF::TPUExecuteOp, TF::TPUExecuteAndUpdateVariablesOp>(
            op.GetBody().front()))
      execute_launches.push_back(op);

      new_op_mapper.map(result, new_result);
    }
  }

  call_op.erase();
}

void UnwrapXlaCallModuleOpPass::runOnOperation() {
  ModuleOp module_op = getOperation();
  SymbolTable symbol_table(module_op);

  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    Region& function_body = func_op.getBody();

    function_body.walk([&](TF::XlaCallModuleOp call_op) {

    if (!isa<TF::IdentityOp, TF::ReadVariableOp>(op)) break;
    v = op->getOperand(0);
  }
  return v;
}

void AnnotateParameterReplicationPass::runOnOperation() {
  ModuleOp m = getOperation();
  OpBuilder builder(m.getContext());
  m.walk([&](tf_device::ClusterFuncOp cluster_func) {
    auto replicate = cluster_func->getParentOfType<tf_device::ReplicateOp>();
    if (!replicate) return;

          return op.getCallTargetName() != "stablehlo.composite";
        });

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

    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns)))) {
      return signalPassFailure();
    }
  }
};

static PassRegistration<LegalizeStablehloCustomCallToCompositePass>

    return op.getOperation();
  }
  auto op = builder.create<tfl::ConstOp>(loc, value);
  op->setAttr("tfl.is_variable", builder.getUnitAttr());
  if (tensor.quantization && !tensor.quantization->min.empty()) {
    if (auto stats_op =
            ConvertMinMaxToStatsOp(tensor, builder, op.getResult())) {
      return stats_op;
    }
  }
  return op.getOperation();
}