include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

// TODO(jpienaar): Move this somewhere general.
class GetI64ScalarElementsAttr<int value> :
  NativeCodeCall<"DenseElementsAttr::get<int64_t>(RankedTensorType::get({}, $_builder.getIntegerType(64)), " # value # ")">;

def FuseMapAndBatch : Pat<
  (TF_BatchDatasetV2Op

  "std::numeric_limits<double>::infinity())">;

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

class GetI64ScalarElementsAttr<int value> :
  NativeCodeCall<"GetI64ElementsAttr({" # value # "}, &$_builder)">;

class GetF32Scalar<int value> :
  NativeCodeCall<"GetF32Scalar(&$_builder, " # value # ")">;