// Unary ops patterns.
//===----------------------------------------------------------------------===//
def IsDataFormatNHWC : ConstantAttr<TF_ConvnetDataFormatAttr, "\"NHWC\"">;
def IsDataFormatNCHW : ConstantAttr<TF_ConvnetDataFormatAttr, "\"NCHW\"">;

class I32VectorElementsAttr<int len> : ElementsAttrBase<
  CPred<"$_self.isa<DenseIntElementsAttr>() &&"

            (Arith_ConstantOp ConstantAttr<I32ElementsAttr, "{0}">)),
          (TF_MatrixSetDiagOp $input, $diagonal)>;

// `align` attribute can be ignored because we only support converting
// `MatrixSetDiagV3` to `MatrixSetDiag` with default `k` inputs.
def ConvertMatrixSetDiagV3ToMatrixSetDiag :
      Pat<(TF_MatrixSetDiagV3Op $input, $diagonal,
            (Arith_ConstantOp ConstantAttr<I32ElementsAttr, "{0}">),

//===----------------------------------------------------------------------===//

class MHLO_ComparisonDirectionValue<string enumStr> :
  ConstantAttr<MHLO_ComparisonDirectionAttr, "::mlir::mhlo::ComparisonDirection::" # enumStr>;

class CHLO_ComparisonDirectionValue<string enumStr> :
  ConstantAttr<CHLO_ComparisonDirectionAttr, "::mlir::chlo::ComparisonDirection::" # enumStr>;

// TODO(b/228291745): Assert that $x and $y have the same shape.

// BiasAddV1 op patterns.
//===----------------------------------------------------------------------===//

def BiasAddV1ToBiasAdd : Pat<
  (TF_BiasAddV1Op:$src $arg0, $arg1),
  (TF_BiasAddOp:$dest $arg0, $arg1,
    ConstantAttr<TF_ConvnetDataFormatAttr, "\"NHWC\"">),
  [], [(CopyAttrs $src, $dest)]>;

//===----------------------------------------------------------------------===//
// Bitcast op patterns.

  [(TF_OneHotOp:$labels $sparse_labels,
     (TF_ConstOp (GetDimSizeAsI32<1> $features, $src_op__0)),
     (TF_ConstOp (GetScalarOfType<1> $features)),
     (TF_ConstOp (GetScalarOfType<0> $features)),
     ConstantAttr<I64Attr, "1">),
   (TF_SelectV2Op:$zero_or_nan
     (TF_LogicalAndOp
       (TF_LessEqualOp
         (TF_ConstOp (GetScalarOfType<0> $sparse_labels)), $sparse_labels),
       (TF_LessOp $sparse_labels,

   HasChloCompareType<"::mlir::chlo::ComparisonType::NOTYPE">]>,
   "compare type supported by TensorFlow">;

class CHLO_ComparisonDirectionValue<string enumStr> :
  ConstantAttr<CHLO_ComparisonDirectionAttr, "::mlir::chlo::ComparisonDirection::" # enumStr>;

class HasMhloCompareType<string value> :
    CPred<"$_self.cast<::mlir::mhlo::ComparisonTypeAttr>().getValue() == " # value>;