UNSORTED_SEGMENT_PROD = 153,
  UNSORTED_SEGMENT_MAX = 154,
  UNSORTED_SEGMENT_SUM = 155,
  ATAN2 = 156,
  UNSORTED_SEGMENT_MIN = 157,
  SIGN = 158,
  BITCAST = 159,
  BITWISE_XOR = 160,
  RIGHT_SHIFT = 161,
  // All Operators start with STABLEHLO_ prefixes are subject to change
  // Many of the ops below can not be executed by TFlite runtime
  STABLEHLO_LOGISTIC = 162, // WARNING: Do not have runtime support

}

// -----

// CHECK-LABEL: testRightShift
func.func @testRightShift(%arg0: tensor<8xui32>, %arg1: tensor<8xui32>) -> tensor<8xui32> {
  // CHECK: "tfl.right_shift"(%arg0, %arg1)
  %0 = "tfl.right_shift"(%arg0, %arg1) : (tensor<8xui32>, tensor<8xui32>) -> tensor<8xui32>
  func.return %0 : tensor<8xui32>
  // CHECK: return %0 : tensor<8xui32>
}

// -----

  %0 = "tf.RightShift"(%arg0, %arg1) : (tensor<8xui32>, tensor<8xui32>) -> tensor<8xui32>
  func.return %0 : tensor<8xui32>

  // CHECK-LABEL: testRightShift
  // CHECK: %[[RES0:.*]] = "tfl.right_shift"(%arg0, %arg1) : (tensor<8xui32>, tensor<8xui32>) -> tensor<8xui32>
  // CHECK: return %[[RES0]] : tensor<8xui32>
}

// =============================================================================
// Training OPs

    TFL_TensorOf<[I8, UI8, I16, UI16, I32, UI32]>:$rhs
  );

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

def TFL_RightShiftOp : TFL_Op<"right_shift", [
    SameOperandsAndResultElementType,
    Pure]> {
  let summary = "Right Shift operator";

  let description = [{
    Elementwise computes the bitwise right-shift of `lhs` by `rhs`.
  }];

for dtype in dtype_list:
  lhs = tf.constant([-1, -5, -3, -14], dtype=dtype)
  rhs = tf.constant([5, 0, 7, 11], dtype=dtype)

  right_shift_result = bitwise_ops.right_shift(lhs, rhs)

  print(right_shift_result)

# This will print:
# tf.Tensor([-1 -5 -1 -1], shape=(4,), dtype=int8)
# tf.Tensor([-1 -5 -1 -1], shape=(4,), dtype=int16)

    "UNSORTED_SEGMENT_PROD",
    "UNSORTED_SEGMENT_MAX",
    "UNSORTED_SEGMENT_SUM",
    "ATAN2",
    "UNSORTED_SEGMENT_MIN",
    "SIGN",
    "BITCAST",
    "BITWISE_XOR",
    "RIGHT_SHIFT",
    "STABLEHLO_LOGISTIC",
    "STABLEHLO_ADD",
    "STABLEHLO_DIVIDE",
    "STABLEHLO_MULTIPLY",
    "STABLEHLO_MAXIMUM",
    "STABLEHLO_RESHAPE",
    "STABLEHLO_CLAMP",

    *   Added int16 indices support for built-in op `gather` and `gather_nd`.
    *   Added 8-bit/16-bit/32-bit int/uint support for the built-in op `right_shift`
    *   Added reference implementation for 16-bit int unquantized `add`.
    *   Added reference implementation for 16-bit int and 32-bit unsigned int unquantized `mul`.