// Ensure bucketize roundtrip exactly

func.func @main(%arg0: tensor<3x2xf32>) -> tensor<3x2xi32> {
  // CHECK-LABEL: @main
  // CHECK: "tfl.bucketize"(%arg0) <{boundaries = [0.000000e+00 : f32, 1.000000e+01 : f32, 1.000000e+02 : f32]}> : (tensor<3x2xf32>) -> tensor<3x2xi32>

  %0 = "tfl.pseudo_const" () {value = dense<[[-5.0, 10000.0], [150.0, 10.0], [5.0, 100.0]]> : tensor<3x2xf32>} : () -> tensor<3x2xf32> loc("Const")
  %1 = "tfl.bucketize"(%0) {boundaries = [0.0 : f32, 10.0 : f32, 100.0 : f32]} : (tensor<3x2xf32>) -> tensor<3x2xi32> loc("bucketize")
  func.return %1 : tensor<3x2xi32>

    func.return %1 : tensor<4x7xcomplex<f64>>
  }

  // CHECK-LABEL: bucketize
  func.func @bucketize(%arg0: tensor<2x5xf32>) -> tensor<2x5xi32> {
    // CHECK-NOT: tf.Bucketize
    %0 = "tf.Bucketize"(%arg0) {boundaries = [0.000000e+00 : f32, 3.000000e+00 : f32, 8.000000e+00 : f32, 1.100000e+01 : f32]} : (tensor<2x5xf32>) -> tensor<2x5xi32>

  // CHECK: return %[[SELECT_V2]] : tensor<8x7x6x5x?x3x2x1xf32>
}

func.func @Bucketize(%arg0: tensor<3x2xf32>) -> tensor<3x2xi32> {
  %0 = "tf.Bucketize"(%arg0) {boundaries = [1.0 : f32, 10.0 : f32, 100.0 : f32]} : (tensor<3x2xf32>) -> tensor<3x2xi32>
  func.return %0: tensor<3x2xi32>

// CHECK-LABEL: Bucketize

  HASHTABLE_SIZE = 139,
  REDUCE_ALL = 140,
  CONV_3D_TRANSPOSE = 141,
  VAR_HANDLE = 142,
  READ_VARIABLE = 143,
  ASSIGN_VARIABLE = 144,
  BROADCAST_ARGS = 145,
  RANDOM_STANDARD_NORMAL = 146,
  BUCKETIZE = 147,
  RANDOM_UNIFORM = 148,
  MULTINOMIAL = 149,
  GELU = 150,
  DYNAMIC_UPDATE_SLICE = 151,
  RELU_0_TO_1 = 152,
  UNSORTED_SEGMENT_PROD = 153,
  UNSORTED_SEGMENT_MAX = 154,

      "BatchMatMulV2",
      "BatchMatMulV3",
      "BatchToSpace",
      "BatchToSpaceND",
      "BesselI0e",
      "BesselI1e",
      "Betainc",
      "BiasAddV1",
      "Bincount",
      "Bucketize",
      "Case",
      "CheckNumerics",
      "Cholesky",
      "ControlTrigger",
      "Conv",
      "Conv2D",
      "Conv2DBackpropFilter",
      "Conv2DBackpropInput",
      "Conv3D",

  let arguments = (ins
    TFL_I32OrI64Tensor:$s0,
    TFL_I32OrI64Tensor:$s1
  );

  let results = (outs
    TFL_I32OrI64Tensor:$r0
  );
}

def TFL_BucketizeOp
    : TFL_Op<"bucketize", [Pure, SameOperandsAndResultShape]> {
  let summary = "Bucketizes 'input' based on 'boundaries'.";

  let description = [{
Example:

If the inputs are `boundaries = [0, 10, 100]` and

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_BoostedTreesBucketizeOp : TF_Op<"BoostedTreesBucketize", [Pure, SameVariadicOperandSize]> {
  let summary = "Bucketize each feature based on bucket boundaries.";

  let description = [{
An op that returns a list of float tensors, where each tensor represents the
bucketized values for a single feature.
  }];

    "HASHTABLE_SIZE",
    "REDUCE_ALL",
    "CONV_3D_TRANSPOSE",
    "VAR_HANDLE",
    "READ_VARIABLE",
    "ASSIGN_VARIABLE",
    "BROADCAST_ARGS",
    "RANDOM_STANDARD_NORMAL",
    "BUCKETIZE",
    "RANDOM_UNIFORM",
    "MULTINOMIAL",
    "GELU",
    "DYNAMIC_UPDATE_SLICE",
    "RELU_0_TO_1",
    "UNSORTED_SEGMENT_PROD",
    "UNSORTED_SEGMENT_MAX",
    "UNSORTED_SEGMENT_SUM",

# Release 2.8.0

## Major Features and Improvements

*   `tf.lite`:

    *   Added TFLite builtin op support for the following TF ops:
        *   `tf.raw_ops.Bucketize` op on CPU.
        *   `tf.where` op for data types
            `tf.int32`/`tf.uint32`/`tf.int8`/`tf.uint8`/`tf.int64`.
        *   `tf.random.normal` op for output data type `tf.float32` on CPU.