outputs=['output: T'])
def _composite_mirror_pad(input_, paddings, mode):
  shape = input_.shape.as_list()
  for i in range(len(shape)):
    rdims = tf.raw_ops.OneHot(
        indices=i, depth=len(shape), on_value=True, off_value=False, axis=-1)
    rarray = tf.raw_ops.Reverse(tensor=input_, dims=rdims)

    left_padding_size = tf.raw_ops.GatherNd(params=paddings, indices=[i, 0])

    "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,

def _tfr_tensor_two_output(x):
  z = array_ops.Split(axis=0, value=x, num_split=2)
  return z[0], z[1]


@composite.Composite('TestNumAttrsOp')
def _tfr_tensor_tensor_with_cst(x1, y1, x2, y2):
  x = array_ops.OneHot(
      indices=[0, 2, -1, x1], depth=y1, on_value=True, off_value=False)
  (x, x2, y2)  # pylint: disable=pointless-statement
  return

#--- test fn for scf control flow ---

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.