// 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

  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