func.return %0#0 : !tfr.tensor
}

// -----

// CHECK-LABEL: quant_rescale
func.func @quant_rescale(%arg0: !tfr.tensor, %arg1: !tfr.tensor, %arg2: i64) -> !tfr.tensor {
  %0 = "tfr.quant_rescale"(%arg0, %arg1, %arg2) : (!tfr.tensor, !tfr.tensor, i64) -> (!tfr.tensor)
  func.return %0 : !tfr.tensor
}

// -----

// CHECK-LABEL: quant_raw_data

  }];
}

def TFR_TFRQuantRescaleOp : TFR_Op<"quant_rescale", [Pure]> {
  let description = [{
   The `quant_rescale` rescales the elements of the integer tensor by the
   floating-point rescale factor. This op needs to be legalized to the preferred
   operations of the backends.

    Example:

    ```mlir
    %3 = tfr.quant_rescale(%2, %1, %0) :
        (tfr.tensor, tfr.tensor, i64) -> (tfr.tensor)

  %cst = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
  %scale_factor = "tfr.cast"(%cst) : (tensor<f32>) -> !tfr.tensor
  %input = "tfr.cast"(%arg0) : (tensor<2xi32>) -> !tfr.tensor
  %rescaled = "tfr.quant_rescale"(%input, %scale_factor, %zp) : (!tfr.tensor, !tfr.tensor, i64) -> !tfr.tensor
  func.return %rescaled : !tfr.tensor

// CHECK-DAG: %[[f32:.*]] = tfr.constant f32 -> !tfr.attr

    return success();
  }
};

class RemoveRescaleOp : public OpRewritePattern<TFRQuantRescaleOp> {
  using OpRewritePattern<TFRQuantRescaleOp>::OpRewritePattern;

 public:
  // Replace quant_rescale (input, scale, zp) with
  // tf.Cast(tf.Round(tf.Cast(input, f32) * scale) + tf.Cast(zp, f32), i32)
  LogicalResult matchAndRewrite(TFRQuantRescaleOp rescale_op,

      CHECK:        %[[act_range:.*]]:2 = tfr.quant_act_range(%{{.*}}, %{{.*}}, %{{.*}}) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
      CHECK:        %[[rescale:.*]] = tfr.quant_rescale(%[[Sub]], %[[factor1]], %{{.*}}) : (!tfr.tensor, !tfr.tensor, i64) -> (!tfr.tensor)
      CHECK:        %[[attr:.*]] = tfr.constant i16 -> !tfr.attr