"MatrixDiagPart",
    "MatrixDiagPartV2",
    "MatrixDiagPartV3",
    "MatrixDiagV2",
    "MatrixDiagV3",
    "MatrixSetDiag",
    "MatrixSetDiagV2",
    "MatrixSetDiagV3",
    "MirrorPad",
    "OneHot",
    "OnesLike",
    "Pack",
    "Pad",
    "PadV2",
    "ParallelConcat",
    "PlaceholderV2",
    "PlaceholderWithDefault",
    "PreventGradient",
    "QuantizeAndDequantize",

  "32 bit integer tensor">;

def ConvertSingleElementAttrToFloatAttr :
  NativeCodeCall<"ConvertSingleElementAttrToFloatAttr($0)">;

// Replace
//   (float)OneHot(index, depth, on_val, off_val, axis)
// With
//   OneHot(index, depth, (float)on_val, (float)off_val, axis)
def FuseOneHotAndCastToFloat : Pat<
  (TFL_CastOp:$output (TFL_OneHotOp $indices,
                                    $depth,

REGISTER_NO_GRADIENT_OP("Shape");
REGISTER_NO_GRADIENT_OP("ShapeN");
REGISTER_NO_GRADIENT_OP("Rank");
REGISTER_NO_GRADIENT_OP("Size");
REGISTER_NO_GRADIENT_OP("BroadcastGradientArgs");
REGISTER_NO_GRADIENT_OP("OneHot");

Status PackGrad(const Scope& scope, const Operation& op,
                const std::vector<Output>& grad_inputs,
                std::vector<Output>* grad_outputs) {
  int N;

  func.return %0 : tensor<?x?xf32>
}

// -----

func.func @testOneHot(%indices: tensor<3xi32>, %depth: tensor<i32>, %on_value: tensor<f32>, %off_value: tensor<f32>) -> tensor<3x5xf32> {
  %result = "tf.OneHot"(%indices, %depth, %on_value, %off_value) {axis = -1 : i64} : (tensor<3xi32>, tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<3x5xf32>
  func.return %result : tensor<3x5xf32>
}

// -----

  // CHECK: return %0 : tensor<2x3x4x5xf32>
}

func.func @OneHot(%arg0: tensor<3xi32>, %arg1: tensor<i32>, %arg2: tensor<f32>, %arg3: tensor<f32>) -> tensor<*xf32> {
  %0 = "tf.OneHot"(%arg0, %arg1, %arg2, %arg3) {axis = -1 : i64} : (tensor<3xi32>, tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
  func.return %0: tensor<*xf32>

// CHECK-LABEL: OneHot

  return true;
}

// Returns true if attr is a DenseIntElementsAttr of int32 or int64 values or an
// incrementing sequence from 0 to N-1.
//
// If such a value is used in an Equal operator, it can be replaced with OneHot.
bool IsOneHotIndexAttribute(Attribute attr) {
  const auto dense_attr = mlir::dyn_cast_or_null<DenseIntElementsAttr>(attr);
  if (!dense_attr) {
    return false;
  }
  auto index_type = dense_attr.getType();

      "Range", "Rank", "Reshape", "Shape", "ShapeN", "Size", "Squeeze",
      "Transpose", "ZerosLike", "OnesLike", "BiasAdd" /*PW + Broadcast*/,
      "BroadcastArgs", "BroadcastGradientArgs", "OneHot", "Concat", "ConcatV2",
      "ConcatOffset", "Const", "MirrorPad", "MirrorPadGrad", "Pack", "Pad",
      "PadV2", "Reverse", "ReverseV2", "ReverseSequence", "Slice", "Split",
      "SplitV", "StridedSlice", "StridedSliceGrad",

  // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
  // CHECK-DAG: %[[LABELS:.*]] = "tf.OneHot"(%[[SPARSE_LABELS]], %[[DEPTH]], %[[ONE]], %[[ZERO]]) <{axis = 1 : i64}> : (tensor<2xi32>, tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<2x3xf32>

  // Adjust labels to have Nan for out of range labels.

*   Fixes a division by zero in TFLite's implementation of `Split`
    ([CVE-2021-29599](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-29599))
*   Fixes a division by zero in TFLite's implementation of `OneHot`
    ([CVE-2021-29600](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-29600))
*   Fixes a division by zero in TFLite's implementation of `DepthwiseConv`

  let hasOptions = 1;
  let customOption = "ReducerOptions";
}

def TFL_OneHotOp : TFL_Op<"one_hot", [
    QuantizableResult,
    Pure]> {
  let summary = "OneHot operator";

  let description = [{
    Returns a one-hot tensor.The locations represented by indices in `indices`
    take value `on_value`, while all other locations take value `off_value`.