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

std::unique_ptr<OperationPass<func::FuncOp>>

};

// Folds tf.IdentityOp and tf.IdentityNOp if op device and the argument devices
// from the defining ops match.
void TensorDeviceCopyConversionPass::runOnOperation() {
  func::FuncOp func_op = getOperation();

  auto should_fold_op_func = [&func_op](const Value &arg,
                                        const StringAttr &op_device) {
    // In TFRT TPU, tensor transfer is handled specifically by D2H and

  if (!quantization_specs.ok()) {
    signalPassFailure();
    return;
  }

  pass_manager.addPass(
      CreateLiftQuantizableSpotsAsFunctionsPass(*quantization_specs));

  if (failed(pass_manager.run(getOperation()))) {
    signalPassFailure();
  }
}

}  // namespace

    : 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());

    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,

}

bool IsMhloUniformQuantizedOp(Operation& op) {
  return llvm::isa<mhlo::UniformQuantizeOp, mhlo::UniformDequantizeOp>(op);
}

void VerifyQuantLegalization::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    // Verify all uq and qint types are lowered.
    if (llvm::any_of(op->getOperandTypes(), IsQuantType) ||

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

    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>