}

func.func @resize_nearest_neighbor(%arg0: tensor<1x100x100x3xf32>, %arg1: tensor<4xi32>) -> tensor<?xf32> {
  %0 = "tf.ResizeNearestNeighbor"(%arg0, %arg1) {align_corners = true} : (tensor<1x100x100x3xf32>, tensor<4xi32>) -> tensor<?xf32>
  func.return %0 : tensor<?xf32>
  // CHECK-LABEL: resize_nearest_neighbor

  //
  // To optimize JAX resize implementation, especially for 2D upscaling, this
  // pattern matching logic captures the following pattern to replace with the
  // single TFL::resize_nearest_neighbor op instance as a fast fused op
  // available in TFLite.
  //
  // - tfl.gather_nd -> tfl.transpose -> tfl.gather_nd -> tfl.transpose
  //   where ...

}

def TFL_ResizeNearestNeighborOp : TFL_Op<"resize_nearest_neighbor", [
    Pure,
    QuantizableResult,
    PredOpTrait<"input and output must have same element type",
      TFL_TCresVTEtIsSameAsOp<0, 0>>,
    TFL_OperandHasRank<0, 4>,
    TFL_OperandHasRank<1, 1>,
    SameOperandsAndResultsScale]> {
  let summary = "ResizeNearestNeighbor Op";

  let description = [{

^bb0(%arg0: tensor<? x ? x ? x ? x !quant.uniform<u8:f32, 0.1>>, %arg1: tensor<? x i32>):
  %0 = "tfl.resize_nearest_neighbor"(%arg0, %arg1) { align_corners = false, half_pixel_centers = false } : (tensor<? x ? x ? x ? x !quant.uniform<u8:f32, 0.1>>, tensor<? x i32>) -> tensor<? x !quant.uniform<u8:f32, 0.1>>