// Adds the `tf._noinline = true` attribute to the function if the name
  // matches.
  if (noinline_functions.contains(func_op.getSymName())) {
    func_op->setAttr(kTfNoinlineAttr, builder.getBoolAttr(true));
    LLVM_DEBUG(llvm::dbgs()
               << "Marked tf._noinline = true: " << func_op.getSymName());
  }
}

static PassRegistration<MarkFunctionsNoinlinePass> pass{};

}  // namespace

                     &replicate.getBody())) {
        // Not a replication-invariant operand.
        continue;
      }
      func.setArgAttr(entry.index(), kReplicationAttr,
                      builder.getBoolAttr(true));
    }
  });
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreateAnnotateParameterReplicationPass() {
  return std::make_unique<AnnotateParameterReplicationPass>();
}

  // in TensorFlow and allows lowering to V1 control flow for loop
  // parallelization.
  dataset_while->setAttr("_lower_using_switch_merge",
                         builder.getBoolAttr(true));

  return dataset_while;
}

// Populate the cond of `dataset_while`.  The cond body just returns the
// condition of whether to continue to next iteration.

  OpBuilder builder(&ctx);
  module_op.walk([&builder](Operation* op) {
    if (op->hasAttrOfType<BoolAttr>("is_stateless") &&
        ContainCalibrationStatisticsSaverOp(op)) {
      op->setAttr("is_stateless", builder.getBoolAttr(false));
    }
  });
}

std::unique_ptr<OperationPass<ModuleOp>>
CreateInsertCalibrationStatisticsSaverPass(
    StringRef calibration_data_dir,

  auto matmul_type =
      RankedTensorType::get(matmul_shape, original_type.getElementType());
  Value out = rewriter.create<TF::BatchMatMulV2Op>(
      op.getLoc(), matmul_type, lhs, rhs, rewriter.getBoolAttr(false),
      rewriter.getBoolAttr(false));

  bool out_reshape_need = (reshape_shape.size() != matmul_shape.size() ||
                           original_type.getRank() != matmul_shape.size());

    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, input_type, min_value, max_value, quant_dim, num_bits,

def convertIntAttrTo32Bit : NativeCodeCall<
    "$_builder.getI32IntegerAttr($0.cast<IntegerAttr>().getInt())">;

// Builds a constant bool attribute.
class GetBoolAttr<int value> :
    NativeCodeCall<"$_builder.getBoolAttr(" # value #")">;

// Converts an integer attribute $0 to 64-bit with builder.
def convertIntAttrTo64Bit : NativeCodeCall<
    "$_builder.getI64IntegerAttr($0.cast<IntegerAttr>().getInt())">;

    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, res_type, min_value, max_value, quant_dim, num_bits,

    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, res_type, min_value, max_value, quant_dim, num_bits,

           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
              mlir::Value y) -> mlir::Operation* {
             return opb
                 .create<mlir::TF::EqualOp>(loc, x, y, opb.getBoolAttr(true))
                 .getOperation();
           });

  // mlir::TF::GreaterEqualOp
  py::class_<mlir::TF::GreaterEqualOp>(m, "Tf_GreaterEqualOp")
      .def("create",