class GetScalarOfFloatType<string value> : NativeCodeCall<
  "GetScalarOfType(getElementTypeOrSelf($0)," # value # ")">;

def GetScalarInfOfType : NativeCodeCall<
  "GetScalarOfType(getElementTypeOrSelf($0), "
  "std::numeric_limits<double>::infinity())">;

def GetScalarNanOfType : NativeCodeCall<
  "GetScalarOfType(getElementTypeOrSelf($0), "
  "std::numeric_limits<double>::quiet_NaN())">;

     $l2_shrinkage, $lr_power, $use_locking, $multiply_linear_by_lr),
  [// grad_with_shrinkage = grad + 2 * l2_shrinkage * var
   (TF_ConstOp:$zero (GetScalarOfType<0> $lr)),
   (TF_ConstOp:$one (GetScalarOfType<1> $lr)),
   (TF_ConstOp:$two (GetScalarOfType<2> $lr)),
   (TF_AddV2Op:$grad_with_shrinkage
     $grad,
     (TF_MulOp
       $two,
       (TF_MulOp $l2_shrinkage, (CreateTFReadVariableOp $src_op, $grad, $var))

// TODO(hinsu): Lower quantized types after supporting them in GetScalarOfType.
def : Pat<(TF_ReluOp AnyTensor:$input),
          (CHLO_BroadcastMaxOp
               (MHLO_ConstantOp:$zero (GetScalarOfType<0> $input)), $input,
               (BinBroadcastDimensions $zero, $input)),
          [(TF_IntOrFpTensor $input)]>;

// TODO(hinsu): Lower quantized types after supporting them in GetScalarOfType.
def : Pat<(TF_Relu6Op AnyRankedTensor:$input),

include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"

// Here, the element type can be any integer or float type. But, note that only
// 32 bit integers are supported for the values.
class GetScalarOfType<int value> : NativeCodeCall<
  "GetScalarOfType(getElementTypeOrSelf($0)," # value # ")">;

// Here, the element type can be any integer or float type. But, note that only
// 32 bit integers are supported for the values.

    Value step_size = rewriter.create<ConstOp>(loc, GetScalarOfType(u64, 256));
    Value increment =
        rewriter.create<MulOp>(loc, u64_scalar, step_size, rng_op.getDelta());

    // Increment the counter.
    SmallVector<Value, 4> pack_args;
    RankedTensorType word_u64_type = RankedTensorType::get({}, u64);
    Value zero_u64 = rewriter.create<ConstOp>(loc, GetScalarOfType(u64, 0));

namespace mlir {
namespace mhlo {

ConstantOp GetScalarConstOfType(Type ty, Location loc, int64_t raw_value,
                                OpBuilder* builder) {
  return builder->create<ConstantOp>(loc, hlo::getScalarOfType(ty, raw_value));
}

ConstantOp GetScalarNegZeroOfType(Type ty, Location loc, OpBuilder* builder) {
  return builder->create<ConstantOp>(loc, hlo::getScalarNegZeroOfType(ty));
}

    auto result_type = x_type;
    auto start = rewriter.create<ConstOp>(loc, GetScalarOfType(int_type, 0));
    Value limit = rewriter.create<ConstOp>(
        loc, GetScalarOfType(int_type, x_type.getShape()[0]));
    auto delta = rewriter.create<ConstOp>(loc, GetScalarOfType(int_type, 1));
    // Construct a sequence of numbers [0, 1, ... len(x)-1].
    auto updates =

namespace mlir {
namespace TF {

// Returns int, float or complex DenseElementsAttr with scalar shape with the
// given element type and the integer value.
template <typename T>
DenseElementsAttr GetScalarOfType(Type ty, T raw_value) {
  RankedTensorType scalar_ty = RankedTensorType::get({}, ty);
  if (auto float_ty = mlir::dyn_cast<FloatType>(ty)) {
    FloatAttr attr = FloatAttr::get(float_ty, raw_value);

def MatrixDiagToV3 : Pat<
  (TF_MatrixDiagOp:$src $diag),
  (TF_MatrixDiagV3Op:$dest $diag, (TF_ConstOp (GetI32Attr<0>)),
    (TF_ConstOp (GetI32Attr<-1>)), (TF_ConstOp (GetI32Attr<-1>)),
    (TF_ConstOp (GetScalarOfType<0> $diag)), (GetStrAttr<"RIGHT_LEFT">)),
  [], [(CopyAttrs $src, $dest)]>;

//===----------------------------------------------------------------------===//
// QuantizeAndDequantizeV2 op patterns.