include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mhlo/IR/hlo_ops.td"

def SignedIntTensor : TensorOf<[I1, I8, I16, I32, I64]>;
def UnsignedIntTensor : TensorOf<[UI8, UI16, UI32, UI64]>;

// IEEE compliant floating point tensors.
def IEEEFloatTensor : TensorOf<[F16, F32, F64]>;

//===----------------------------------------------------------------------===//
// BatchNorm op patterns.

  let arguments = (ins
     TFL_TensorOf<[I8, I16, I32, I64, UI8, UI16, UI32, UI64, F32, F64]>:$input,
     TFL_TensorOf<[I32]>:$dilations,
     TFL_0DTensorOf<[I8, I16, I32, I64, UI8, UI16, UI32, UI64, F32, F64]>:$padding_value
  );

  let results = (outs TFL_TensorOf<[I8, I16, I32, I64, UI8, UI16, UI32, UI64, F32, F64]>:$output);
}

def TFL_AddOp : TFL_Op<"add", [

  TF_DerivedResultTypeAttr dtype = TF_DerivedResultTypeAttr<0>;
}

// TODO(lyandy): Investigate supported dtypes (`minval`, `maxval`, `output`) for
// `tf.StatefulUniformInt`. tf2xla kernels support i32, i64, ui32, and ui64
// while TensorFlow CPU/GPU kernels only support i32 and i64.
def TF_StatefulUniformIntOp : TF_Op<"StatefulUniformInt", []> {
  let summary = "Outputs random integers from a uniform distribution.";

// CHECK:           %[[CST:.*]] = "tf.Const"() <{value = dense<18446744073709551615> : tensor<ui64>}> : () -> tensor<ui64>
// CHECK:           %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xui64>, tensor<ui64>) -> tensor<7x9x11xui64>
// CHECK:           return %[[RES]] : tensor<7x9x11xui64>
// CHECK:         }

    if (!type || type.getElementType().isUnsignedInteger(64)) {
      return rewriter.notifyMatchFailure(dynamic_iota_op,
                                         "TF::RangeOp doesn't support UI64");
    }
    // Only support 1D for now.
    if (type.getRank() > 1 || dynamic_iota_op.getIotaDimension() != 0) {
      return rewriter.notifyMatchFailure(
          dynamic_iota_op, [&](::mlir::Diagnostic& diag) {