mlir::TF::AssignVariableOp assign_var_op, BundleWriter& bundle_writer) {
  auto resource_operand = assign_var_op.getOperand(0);
  auto var_handle_op =
      llvm::dyn_cast<mlir::TF::VarHandleOp>(resource_operand.getDefiningOp());
  if (!var_handle_op) {
    assign_var_op->emitRemark(
        "Operand idx 0 is not a tf.VarHandleOp. The initializing tensor is not "
        "saved to checkpoint.");
    return "";
  }

                             Value multiplier) {
  auto dq_op = value.getDefiningOp<quantfork::DequantizeCastOp>();
  if (!dq_op) {
    auto mul_op = builder.create<TF::MulOp>(loc, value, multiplier);
    return mul_op.getResult();
  }
  auto q_op = dq_op.getArg().getDefiningOp<quantfork::QuantizeCastOp>();
  if (!q_op) return {};

  Value float_value = q_op.getArg();

    // Operand must be defined by a transpose op.
    TransposeOp transpose =
        dyn_cast_or_null<TransposeOp>(operand.get().getDefiningOp());
    if (!transpose) return;

    // With permutation defined by constant operation.
    ConstOp perm =
        dyn_cast_or_null<ConstOp>(transpose.getOperand(1).getDefiningOp());
    if (!perm) return;

    // With the same permutation indices.

    for (const auto& operand : same_scale_op->getOperands()) {
      auto dq_op = dyn_cast_or_null<quantfork::DequantizeCastOp>(
          operand.getDefiningOp());
      if (!dq_op) continue;

      Operation* preceding_op = dq_op.getArg().getDefiningOp();
      if (!preceding_op) continue;

      // Check whether the preceding op is a quantized composite function.

  LogicalResult matchAndRewrite(quantfork::DequantizeCastOp op,
                                PatternRewriter& rewriter) const override {
    auto input_op = op.getArg().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<quantfork::QuantizeCastOp>(input_op)) {
      if (!q->getAttr(kVolatileOpAttrName)) return failure();

      if (remove_volatile_ops_type == kPreserveInputsAndOutputs) {

def CloneF32ElementsAttrWithOnes
  : NativeCodeCall<"DenseElementsAttr::get($0.getType().cast<ShapedType>(), (float)1.0)">;

def NotConstant : Constraint<
    CPred<"$0.isa<BlockArgument>() || !llvm::isa<stablehlo::ConstantOp>($0.getDefiningOp())">,
    "Is not a constant.">;

def : Pat<(StableHLO_DivOp $l,
            (StableHLO_ConstantOp:$divisor FloatElementsAttr<32>:$cst)),
          (StableHLO_MulOp $l,
            (StableHLO_DivOp

                            (CreateTFLCastToInt32Op (TFL_ConstOp $perm)))>;


def ConvertDotGeneralOp : NativeCodeCall<"ConvertDotGeneralOp($_builder, "
                                               "$0.getDefiningOp())">;
def LegalizeDotGeneral: Pat<(MHLO_DotGeneralOp:$old_value
               $lhs,
               $rhs,
               $dot_dimension_numbers, $precision_config),

  Operation* xla_call_module_op =
      FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
  Operation* filter_dcast_op =
      xla_call_module_op->getOperand(1).getDefiningOp();
  Operation* filter_qcast_op = filter_dcast_op->getOperand(0).getDefiningOp();
  ASSERT_NE(filter_qcast_op, nullptr);
  EXPECT_TRUE(isa<quantfork::QuantizeCastOp>(filter_qcast_op));
  EXPECT_TRUE(isa<quantfork::DequantizeCastOp>(filter_dcast_op));

// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

void TPUAnnotateDynamicShapeInputsPass::runOnOperation() {
  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    Builder builder(cluster_func_op->getContext());

// Combines the two variadic arguments ($in_tensors and $captured_tensors).
def GetBatchFunctionOpArgOperands:
    NativeCodeCall<"cast<TF::BatchFunctionOp>($0[0].getDefiningOp()).getArgOperands()">;

// Replaces `TF_BatchFunctionOp` into `TF_PartitionedCallOp` that calls the
// same $f. This may be required, for example, when inlining is desired,