// CHECK:           %[[VAL_0:.*]] = "tf.All"(%[[ARG_0]], %[[DIMENSIONS]]) <{keep_dims = false}> : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1>
// CHECK:           return %[[VAL_0:.*]] : tensor<2x4x5xi1>
// CHECK:         }
func.func @convert_reduce_to_all(%arg0: tensor<1x2x3x4x5xi1>, %arg1: tensor<2xi64>) -> tensor<2x4x5xi1> {
  %0 = mhlo.constant dense<true> : tensor<i1>
  %1 = "mhlo.reduce"(%arg0, %0) ({

    %9 = "tf.Identity"(%cst_2) {device = ""} : (tensor<2x4x5xi8>) -> tensor<2x4x5xi8>
    %10 = "tf.Cast"(%9) {Truncate = false, device = ""} : (tensor<2x4x5xi8>) -> tensor<2x4x5xi32>
    %11 = "tf.Einsum"(%8, %10) {device = "", equation = "abc,acd->abd"} : (tensor<2x3x4xi32>, tensor<2x4x5xi32>) -> tensor<2x3x5xi32>

  // CHECK-DAG: %[[NEG:.*]] = "tf.Neg"(%[[NUM_UPPER_OR_N]]) : (tensor<i64>) -> tensor<i64>
  // CHECK-DAG: %[[GE:.*]] = "tf.GreaterEqual"(%[[OFFSET]], %[[NEG]]) : (tensor<4x5xi64>, tensor<i64>) -> tensor<4x5xi1>
  // CHECK-DAG: %[[INDICATOR:.*]] = "tf.LogicalAnd"(%[[LE2]], %[[GE]]) : (tensor<4x5xi1>, tensor<4x5xi1>) -> tensor<4x5xi1>

func.func @ReplaceReshapeEqualWithOneHotWithBatchingDim(%arg: tensor<2x2x1xi32>) -> tensor<2x2x3xi1> {
  %cst = arith.constant dense<[0, 1, 2]> : tensor<3xi32>
  %result = "tfl.equal"(%arg, %cst) : (tensor<2x2x1xi32>, tensor<3xi32>) -> tensor<2x2x3xi1>
  func.return %result : tensor<2x2x3xi1>

  // CHECK-DAG: %[[CST1:.*]] = arith.constant dense<3> : tensor<i32>

    // CHECK-NEXT: %[[CMP:.*]] = mhlo.compare GT, %[[INP:.*]], %[[ZERO]], NOTYPE : (tensor<1x4x4xf32>, tensor<1x4x4xf32>) -> tensor<1x4x4xi1>
    // CHECK-NEXT: %[[RES:.*]] = mhlo.select %[[CMP]], %[[GRADIENT]], %[[LEAKYGRAD]] : tensor<1x4x4xi1>, tensor<1x4x4xf32>
    // CHECK-NEXT: return %[[RES]] : tensor<1x4x4xf32>