}];
}

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,