Merge branch 'main' into example-reorg

examples/onnx/1l_tiny_div/gen.py

@@ -0,0 +1,39 @@
+from torch import nn
+import torch
+import json
+
+class Circuit(nn.Module):
+    def __init__(self, inplace=False):
+        super(Circuit, self).__init__()
+
+    def forward(self, x):
+        return x/ 10000
+
+
+circuit = Circuit()
+
+
+x = torch.empty(1, 8).random_(0, 2)
+
+out = circuit(x)
+
+print(out)
+
+torch.onnx.export(circuit, x, "network.onnx",
+                  export_params=True,        # store the trained parameter weights inside the model file
+                  opset_version=17,          # the ONNX version to export the model to
+                  do_constant_folding=True,  # whether to execute constant folding for optimization
+                  input_names=['input'],   # the model's input names
+                  output_names=['output'],  # the model's output names
+                  dynamic_axes={'input': {0: 'batch_size'},  # variable length axes
+                                'output': {0: 'batch_size'}})
+
+
+d1 = ((x).detach().numpy()).reshape([-1]).tolist()
+
+data = dict(
+    input_data=[d1],
+)
+
+# Serialize data into file:
+json.dump(data, open("input.json", 'w'))

examples/onnx/1l_tiny_div/input.json

@@ -0,0 +1 @@
+{"input_data": [[1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0]]}

src/circuit/ops/hybrid.rs

@@ -1,7 +1,8 @@
 use super::*;
 use crate::{
-    circuit::{self, layouts, utils, Tolerance},
+    circuit::{layouts, utils, Tolerance},
     fieldutils::{felt_to_i128, i128_to_felt},
+    graph::multiplier_to_scale,
     tensor::{self, Tensor, TensorError, TensorType, ValTensor},
 };
 use halo2curves::ff::PrimeField;
@@ -13,6 +14,15 @@ use serde::{Deserialize, Serialize};
 /// An enum representing the operations that consist of both lookups and arithmetic operations.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub enum HybridOp {
+    Recip {
+        input_scale: utils::F32,
+        output_scale: utils::F32,
+        use_range_check_for_int: bool,
+    },
+    Div {
+        denom: utils::F32,
+        use_range_check_for_int: bool,
+    },
     ReduceMax {
         axes: Vec<usize>,
     },
@@ -75,6 +85,7 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
         match self {
             HybridOp::Greater | HybridOp::Less | HybridOp::Equals => vec![0, 1],
             HybridOp::ScatterElements { .. } => vec![0, 2],
+            HybridOp::GreaterEqual | HybridOp::LessEqual => vec![0, 1],
             _ => vec![],
         }
     }
@@ -113,6 +124,40 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
                     - tensor::ops::sum(&tensor::ops::nonlinearities::leakyrelu(&inter_1, 0.0))?)?;
                 (res.clone(), vec![inter_1, inter_2])
             }
+            HybridOp::Div {
+                denom,
+                use_range_check_for_int,
+                ..
+            } => {
+                let res = crate::tensor::ops::nonlinearities::const_div(&x, denom.0 as f64);
+                // if denom is a round number and use_range_check_for_int is true, use range check check
+                if denom.0.fract() == 0.0 && *use_range_check_for_int {
+                    let divisor = Tensor::from(vec![denom.0 as i128 / 2].into_iter());
+                    (res, vec![-divisor.clone(), divisor])
+                } else {
+                    (res, vec![x])
+                }
+            }
+            HybridOp::Recip {
+                input_scale,
+                output_scale,
+                use_range_check_for_int,
+            } => {
+                let res = crate::tensor::ops::nonlinearities::recip(
+                    &x,
+                    input_scale.0 as f64,
+                    output_scale.0 as f64,
+                );
+                // if scale is a round number and use_range_check_for_int is true, use range check check
+                if input_scale.0.fract() == 0.0 && *use_range_check_for_int {
+                    let err_tol = Tensor::from(
+                        vec![(output_scale.0 * input_scale.0) as i128 / 2].into_iter(),
+                    );
+                    (res, vec![-err_tol.clone(), err_tol])
+                } else {
+                    (res, vec![x])
+                }
+            }
             HybridOp::ReduceArgMax { dim } => {
                 let res = tensor::ops::argmax_axes(&x, *dim)?;
                 let indices = Tensor::from(0..x.dims()[*dim] as i128);
@@ -272,6 +317,21 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
 
     fn as_string(&self) -> String {
         match self {
+            HybridOp::Recip {
+                input_scale,
+                output_scale,
+                use_range_check_for_int,
+            } => format!(
+                "RECIP (input_scale={}, output_scale={}, use_range_check_for_int={})",
+                input_scale, output_scale, use_range_check_for_int
+            ),
+            HybridOp::Div {
+                denom,
+                use_range_check_for_int,
+            } => format!(
+                "DIV (denom={}, use_range_check_for_int={})",
+                denom, use_range_check_for_int
+            ),
             HybridOp::SumPool {
                 padding,
                 stride,
@@ -335,6 +395,57 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
                 *kernel_shape,
                 *normalized,
             )?,
+            HybridOp::Recip {
+                input_scale,
+                output_scale,
+                use_range_check_for_int,
+            } => {
+                if input_scale.0.fract() == 0.0
+                    && output_scale.0.fract() == 0.0
+                    && *use_range_check_for_int
+                {
+                    layouts::recip(
+                        config,
+                        region,
+                        values[..].try_into()?,
+                        i128_to_felt(input_scale.0 as i128),
+                        i128_to_felt(output_scale.0 as i128),
+                    )?
+                } else {
+                    layouts::nonlinearity(
+                        config,
+                        region,
+                        values.try_into()?,
+                        &LookupOp::Recip {
+                            input_scale: *input_scale,
+                            output_scale: *output_scale,
+                        },
+                    )?
+                }
+            }
+            HybridOp::Div {
+                denom,
+                use_range_check_for_int,
+                ..
+            } => {
+                if denom.0.fract() == 0.0 && *use_range_check_for_int {
+                    layouts::div(
+                        config,
+                        region,
+                        values[..].try_into()?,
+                        i128_to_felt(denom.0 as i128),
+                    )?
+                } else {
+                    layouts::nonlinearity(
+                        config,
+                        region,
+                        values.try_into()?,
+                        &LookupOp::Div {
+                            denom: denom.clone(),
+                        },
+                    )?
+                }
+            }
             HybridOp::Gather { dim, constant_idx } => {
                 if let Some(idx) = constant_idx {
                     tensor::ops::gather(values[0].get_inner_tensor()?, idx, *dim)?.into()
@@ -422,86 +533,12 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
             | HybridOp::OneHot { .. }
             | HybridOp::ReduceArgMin { .. } => 0,
             HybridOp::Softmax { .. } => 2 * in_scales[0],
+            HybridOp::Recip { output_scale, .. } => multiplier_to_scale(output_scale.0 as f64),
             _ => in_scales[0],
         };
         Ok(scale)
     }
 
-    fn required_lookups(&self) -> Vec<LookupOp> {
-        match self {
-            HybridOp::ReduceMax { .. }
-            | HybridOp::ReduceMin { .. }
-            | HybridOp::MaxPool2d { .. } => Op::<F>::required_lookups(&LookupOp::ReLU),
-            HybridOp::Softmax { scale, .. } => {
-                vec![
-                    LookupOp::Exp { scale: *scale },
-                    LookupOp::Recip {
-                        scale: scale.0.powf(2.0).into(),
-                    },
-                ]
-            }
-            HybridOp::RangeCheck(tol) => {
-                let mut lookups = vec![];
-                if tol.val > 0.0 {
-                    let scale_squared = tol.scale.0.powf(2.0);
-                    lookups.extend([
-                        LookupOp::Recip {
-                            scale: scale_squared.into(),
-                        },
-                        LookupOp::GreaterThan {
-                            a: circuit::utils::F32((tol.val * scale_squared) / 100.0),
-                        },
-                    ]);
-                }
-                lookups
-            }
-            HybridOp::Greater { .. } | HybridOp::Less { .. } => {
-                vec![LookupOp::GreaterThan {
-                    a: circuit::utils::F32(0.),
-                }]
-            }
-            HybridOp::GreaterEqual { .. } | HybridOp::LessEqual { .. } => {
-                vec![LookupOp::GreaterThanEqual {
-                    a: circuit::utils::F32(0.),
-                }]
-            }
-            HybridOp::TopK { .. } => {
-                vec![
-                    LookupOp::GreaterThan {
-                        a: circuit::utils::F32(0.),
-                    },
-                    LookupOp::KroneckerDelta,
-                ]
-            }
-            HybridOp::Gather {
-                constant_idx: None, ..
-            }
-            | HybridOp::OneHot { .. }
-            | HybridOp::GatherElements {
-                constant_idx: None, ..
-            }
-            | HybridOp::ScatterElements {
-                constant_idx: None, ..
-            }
-            | HybridOp::Equals { .. } => {
-                vec![LookupOp::KroneckerDelta]
-            }
-            HybridOp::ReduceArgMax { .. } | HybridOp::ReduceArgMin { .. } => {
-                vec![LookupOp::ReLU, LookupOp::KroneckerDelta]
-            }
-            HybridOp::SumPool {
-                kernel_shape,
-                normalized: true,
-                ..
-            } => {
-                vec![LookupOp::Div {
-                    denom: utils::F32((kernel_shape.0 * kernel_shape.1) as f32),
-                }]
-            }
-            _ => vec![],
-        }
-    }
-
     fn clone_dyn(&self) -> Box<dyn Op<F>> {
         Box::new(self.clone()) // Forward to the derive(Clone) impl
     }

src/circuit/ops/layouts.rs

@@ -18,7 +18,10 @@ use super::{
     region::RegionCtx,
 };
 use crate::{
-    circuit::{ops::base::BaseOp, utils},
+    circuit::{
+        ops::base::BaseOp,
+        utils::{self},
+    },
     fieldutils::{felt_to_i128, i128_to_felt},
     tensor::{
         get_broadcasted_shape,
@@ -61,7 +64,7 @@ pub fn div<F: PrimeField + TensorType + PartialOrd>(
     let input = value[0].clone();
     let input_dims = input.dims();
 
-    let range_check_bracket = felt_to_i128(div) - 1;
+    let range_check_bracket = felt_to_i128(div) / 2;
 
     let mut divisor = Tensor::from(vec![ValType::Constant(div)].into_iter());
     divisor.set_visibility(&crate::graph::Visibility::Fixed);
@@ -72,8 +75,7 @@ pub fn div<F: PrimeField + TensorType + PartialOrd>(
 
     let mut claimed_output: ValTensor<F> = if is_assigned {
         let input_evals = input.get_int_evals()?;
-        let divisor_evals = divisor.get_int_evals()?;
-        tensor::ops::div(&[input_evals.clone(), divisor_evals.clone()])?
+        tensor::ops::nonlinearities::const_div(&input_evals.clone(), felt_to_i128(div) as f64)
             .iter()
             .map(|x| Ok(Value::known(i128_to_felt(*x))))
             .collect::<Result<Tensor<Value<F>>, Box<dyn Error>>>()?
@@ -94,6 +96,8 @@ pub fn div<F: PrimeField + TensorType + PartialOrd>(
         BaseOp::Mult,
     )?;
 
+    log::debug!("product: {:?}", product.get_int_evals()?);
+
     let diff_with_input = pairwise(
         config,
         region,
@@ -111,6 +115,83 @@ pub fn div<F: PrimeField + TensorType + PartialOrd>(
     Ok(claimed_output)
 }
 
+/// recip accumulated layout
+pub fn recip<F: PrimeField + TensorType + PartialOrd>(
+    config: &BaseConfig<F>,
+    region: &mut RegionCtx<F>,
+    value: &[ValTensor<F>; 1],
+    input_scale: F,
+    output_scale: F,
+) -> Result<ValTensor<F>, Box<dyn Error>> {
+    let input = value[0].clone();
+    let input_dims = input.dims();
+
+    let range_check_bracket = felt_to_i128(output_scale * input_scale) / 2;
+
+    let mut scaled_unit =
+        Tensor::from(vec![ValType::Constant(output_scale * input_scale)].into_iter());
+    scaled_unit.set_visibility(&crate::graph::Visibility::Fixed);
+    let scaled_unit = region.assign(&config.inputs[1], &scaled_unit.into())?;
+    region.increment(scaled_unit.len());
+
+    let is_assigned = !input.any_unknowns()? && !scaled_unit.any_unknowns()?;
+
+    let mut claimed_output: ValTensor<F> = if is_assigned {
+        let input_evals = input.get_int_evals()?;
+        tensor::ops::nonlinearities::recip(
+            &input_evals,
+            felt_to_i128(input_scale) as f64,
+            felt_to_i128(output_scale) as f64,
+        )
+        .iter()
+        .map(|x| Ok(Value::known(i128_to_felt(*x))))
+        .collect::<Result<Tensor<Value<F>>, Box<dyn Error>>>()?
+        .into()
+    } else {
+        Tensor::new(
+            Some(&vec![Value::<F>::unknown(); input.len()]),
+            &[input.len()],
+        )?
+        .into()
+    };
+    claimed_output.reshape(input_dims)?;
+
+    // this is now of scale 2 * scale
+    let product = pairwise(
+        config,
+        region,
+        &[claimed_output.clone(), input.clone()],
+        BaseOp::Mult,
+    )?;
+
+    log::debug!("product: {:?}", product.get_int_evals()?);
+
+    // this is now of scale 2 * scale hence why we rescaled the unit scale
+    let diff_with_input = pairwise(
+        config,
+        region,
+        &[product.clone(), scaled_unit.clone()],
+        BaseOp::Sub,
+    )?;
+
+    log::debug!("scaled_unit: {:?}", scaled_unit.get_int_evals()?);
+
+    // debug print the diff
+    log::debug!("diff_with_input: {:?}", diff_with_input.get_int_evals()?);
+
+    log::debug!("range_check_bracket: {:?}", range_check_bracket);
+
+    // at most the error should be in the original unit scale's range
+    range_check(
+        config,
+        region,
+        &[diff_with_input],
+        &(-range_check_bracket, range_check_bracket),
+    )?;
+
+    Ok(claimed_output)
+}
+
 /// Dot product accumulated layout
 pub fn dot<F: PrimeField + TensorType + PartialOrd>(
     config: &BaseConfig<F>,
@@ -2371,6 +2452,8 @@ pub fn range_check<F: PrimeField + TensorType + PartialOrd>(
     values: &[ValTensor<F>; 1],
     range: &crate::circuit::table::Range,
 ) -> Result<ValTensor<F>, Box<dyn Error>> {
+    region.add_used_range_check(*range);
+
     // time the entire operation
     let timer = instant::Instant::now();
 
@@ -2415,6 +2498,8 @@ pub fn nonlinearity<F: PrimeField + TensorType + PartialOrd>(
     values: &[ValTensor<F>; 1],
     nl: &LookupOp,
 ) -> Result<ValTensor<F>, Box<dyn Error>> {
+    region.add_used_lookup(nl.clone());
+
     // time the entire operation
     let timer = instant::Instant::now();
 
@@ -2884,7 +2969,8 @@ pub fn softmax<F: PrimeField + TensorType + PartialOrd>(
         &[denom],
         // we set to input scale + output_scale so the output scale is output)scale
         &LookupOp::Recip {
-            scale: scale.0.powf(2.0).into(),
+            input_scale: scale,
+            output_scale: scale,
         },
     )?;
 
@@ -2912,19 +2998,22 @@ pub fn range_check_percent<F: PrimeField + TensorType + PartialOrd>(
     // Calculate the difference between the expected output and actual output
     let diff = pairwise(config, region, values, BaseOp::Sub)?;
 
-    let scale_squared = scale.0.powf(2.0);
     // Calculate the reciprocal of the expected output tensor, scaling by double the scaling factor
     let recip = nonlinearity(
         config,
         region,
         &[values[0].clone()],
         &LookupOp::Recip {
-            scale: scale_squared.into(),
+            input_scale: scale,
+            output_scale: scale,
         },
     )?;
+
     // Multiply the difference by the recip
     let product = pairwise(config, region, &[diff, recip], BaseOp::Mult)?;
 
+    let scale_squared = scale.0 * scale.0;
+
     // Use the greater than look up table to check if the percent error is within the tolerance for upper bound
     let tol = tol / 100.0;
     let upper_bound = nonlinearity(

src/circuit/ops/lookup.rs

@@ -5,7 +5,7 @@ use std::error::Error;
 use crate::{
     circuit::{layouts, table::Range, utils},
     fieldutils::{felt_to_i128, i128_to_felt},
-    graph::{multiplier_to_scale, scale_to_multiplier},
+    graph::multiplier_to_scale,
     tensor::{self, Tensor, TensorError, TensorType},
 };
 
@@ -17,42 +17,112 @@ use halo2curves::ff::PrimeField;
 #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Deserialize, Serialize)]
 pub enum LookupOp {
     Abs,
-    Div { denom: utils::F32 },
-    Cast { scale: utils::F32 },
+    Div {
+        denom: utils::F32,
+    },
+    Cast {
+        scale: utils::F32,
+    },
     ReLU,
-    Max { scale: utils::F32, a: utils::F32 },
-    Min { scale: utils::F32, a: utils::F32 },
-    Ceil { scale: utils::F32 },
-    Floor { scale: utils::F32 },
-    Round { scale: utils::F32 },
-    RoundHalfToEven { scale: utils::F32 },
-    Sqrt { scale: utils::F32 },
-    Rsqrt { scale: utils::F32 },
-    Recip { scale: utils::F32 },
-    LeakyReLU { slope: utils::F32 },
-    Sigmoid { scale: utils::F32 },
-    Ln { scale: utils::F32 },
-    Exp { scale: utils::F32 },
-    Cos { scale: utils::F32 },
-    ACos { scale: utils::F32 },
-    Cosh { scale: utils::F32 },
-    ACosh { scale: utils::F32 },
-    Sin { scale: utils::F32 },
-    ASin { scale: utils::F32 },
-    Sinh { scale: utils::F32 },
-    ASinh { scale: utils::F32 },
-    Tan { scale: utils::F32 },
-    ATan { scale: utils::F32 },
-    Tanh { scale: utils::F32 },
-    ATanh { scale: utils::F32 },
-    Erf { scale: utils::F32 },
-    GreaterThan { a: utils::F32 },
-    LessThan { a: utils::F32 },
-    GreaterThanEqual { a: utils::F32 },
-    LessThanEqual { a: utils::F32 },
+    Max {
+        scale: utils::F32,
+        a: utils::F32,
+    },
+    Min {
+        scale: utils::F32,
+        a: utils::F32,
+    },
+    Ceil {
+        scale: utils::F32,
+    },
+    Floor {
+        scale: utils::F32,
+    },
+    Round {
+        scale: utils::F32,
+    },
+    RoundHalfToEven {
+        scale: utils::F32,
+    },
+    Sqrt {
+        scale: utils::F32,
+    },
+    Rsqrt {
+        scale: utils::F32,
+    },
+    Recip {
+        input_scale: utils::F32,
+        output_scale: utils::F32,
+    },
+    LeakyReLU {
+        slope: utils::F32,
+    },
+    Sigmoid {
+        scale: utils::F32,
+    },
+    Ln {
+        scale: utils::F32,
+    },
+    Exp {
+        scale: utils::F32,
+    },
+    Cos {
+        scale: utils::F32,
+    },
+    ACos {
+        scale: utils::F32,
+    },
+    Cosh {
+        scale: utils::F32,
+    },
+    ACosh {
+        scale: utils::F32,
+    },
+    Sin {
+        scale: utils::F32,
+    },
+    ASin {
+        scale: utils::F32,
+    },
+    Sinh {
+        scale: utils::F32,
+    },
+    ASinh {
+        scale: utils::F32,
+    },
+    Tan {
+        scale: utils::F32,
+    },
+    ATan {
+        scale: utils::F32,
+    },
+    Tanh {
+        scale: utils::F32,
+    },
+    ATanh {
+        scale: utils::F32,
+    },
+    Erf {
+        scale: utils::F32,
+    },
+    GreaterThan {
+        a: utils::F32,
+    },
+    LessThan {
+        a: utils::F32,
+    },
+    GreaterThanEqual {
+        a: utils::F32,
+    },
+    LessThanEqual {
+        a: utils::F32,
+    },
     Sign,
     KroneckerDelta,
-    Pow { scale: utils::F32, a: utils::F32 },
+    Pow {
+        scale: utils::F32,
+        a: utils::F32,
+    },
 }
 
 impl LookupOp {
@@ -120,7 +190,14 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for LookupOp {
                 &x,
                 f32::from(*scale).into(),
             )),
-            LookupOp::Recip { scale } => Ok(tensor::ops::nonlinearities::recip(&x, scale.into())),
+            LookupOp::Recip {
+                input_scale,
+                output_scale,
+            } => Ok(tensor::ops::nonlinearities::recip(
+                &x,
+                input_scale.into(),
+                output_scale.into(),
+            )),
             LookupOp::ReLU => Ok(tensor::ops::nonlinearities::leakyrelu(&x, 0_f64)),
 
             LookupOp::LeakyReLU { slope: a } => {
@@ -173,7 +250,13 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for LookupOp {
             LookupOp::GreaterThanEqual { .. } => "GREATER_THAN_EQUAL".into(),
             LookupOp::LessThan { .. } => "LESS_THAN".into(),
             LookupOp::LessThanEqual { .. } => "LESS_THAN_EQUAL".into(),
-            LookupOp::Recip { scale, .. } => format!("RECIP(scale={})", scale),
+            LookupOp::Recip {
+                input_scale,
+                output_scale,
+            } => format!(
+                "RECIP(input_scale={}, output_scale={})",
+                input_scale, output_scale
+            ),
             LookupOp::Div { denom, .. } => format!("DIV(denom={})", denom),
             LookupOp::Cast { scale } => format!("CAST(scale={})", scale),
             LookupOp::Ln { scale } => format!("LN(scale={})", scale),
@@ -220,12 +303,7 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for LookupOp {
                 let in_scale = inputs_scale[0];
                 in_scale + multiplier_to_scale(1. / scale.0 as f64)
             }
-            LookupOp::Recip { scale } => {
-                let mut out_scale = inputs_scale[0];
-                out_scale +=
-                    multiplier_to_scale(scale.0 as f64 / scale_to_multiplier(out_scale).powf(2.0));
-                out_scale
-            }
+            LookupOp::Recip { output_scale, .. } => multiplier_to_scale(output_scale.into()),
             LookupOp::Sign
             | LookupOp::GreaterThan { .. }
             | LookupOp::LessThan { .. }
@@ -237,10 +315,6 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for LookupOp {
         Ok(scale)
     }
 
-    fn required_lookups(&self) -> Vec<LookupOp> {
-        vec![self.clone()]
-    }
-
     fn clone_dyn(&self) -> Box<dyn Op<F>> {
         Box::new(self.clone()) // Forward to the derive(Clone) impl
     }

src/circuit/ops/mod.rs

@@ -10,8 +10,6 @@ use halo2curves::ff::PrimeField;
 
 use self::{lookup::LookupOp, region::RegionCtx};
 
-use super::table::Range;
-
 ///
 pub mod base;
 ///
@@ -57,16 +55,6 @@ pub trait Op<F: PrimeField + TensorType + PartialOrd>: std::fmt::Debug + Send +
         vec![]
     }
 
-    /// Returns the lookups required by the operation.
-    fn required_lookups(&self) -> Vec<LookupOp> {
-        vec![]
-    }
-
-    /// Returns the range checks required by the operation.
-    fn required_range_checks(&self) -> Vec<Range> {
-        vec![]
-    }
-
     /// Returns true if the operation is an input.
     fn is_input(&self) -> bool {
         false

src/circuit/ops/poly.rs

@@ -33,7 +33,9 @@ pub enum PolyOp {
     Sub,
     Neg,
     Mult,
-    Identity,
+    Identity {
+        out_scale: Option<crate::Scale>,
+    },
     Reshape(Vec<usize>),
     MoveAxis {
         source: usize,
@@ -85,7 +87,9 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<
             PolyOp::Resize { .. } => "RESIZE".into(),
             PolyOp::Iff => "IFF".into(),
             PolyOp::Einsum { equation, .. } => format!("EINSUM {}", equation),
-            PolyOp::Identity => "IDENTITY".into(),
+            PolyOp::Identity { out_scale } => {
+                format!("IDENTITY (out_scale={:?})", out_scale)
+            }
             PolyOp::Reshape(shape) => format!("RESHAPE (shape={:?})", shape),
             PolyOp::Flatten(_) => "FLATTEN".into(),
             PolyOp::Pad(_) => "PAD".into(),
@@ -135,7 +139,7 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<
             PolyOp::Resize { scale_factor } => tensor::ops::resize(&inputs[0], scale_factor),
             PolyOp::Iff => tensor::ops::iff(&inputs[0], &inputs[1], &inputs[2]),
             PolyOp::Einsum { equation } => tensor::ops::einsum(equation, &inputs),
-            PolyOp::Identity => Ok(inputs[0].clone()),
+            PolyOp::Identity { .. } => Ok(inputs[0].clone()),
             PolyOp::Reshape(new_dims) => {
                 let mut t = inputs[0].clone();
                 t.reshape(new_dims)?;
@@ -264,7 +268,7 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<
             PolyOp::Mult => {
                 layouts::pairwise(config, region, values[..].try_into()?, BaseOp::Mult)?
             }
-            PolyOp::Identity => layouts::identity(config, region, values[..].try_into()?)?,
+            PolyOp::Identity { .. } => layouts::identity(config, region, values[..].try_into()?)?,
             PolyOp::Reshape(d) | PolyOp::Flatten(d) => layouts::reshape(values[..].try_into()?, d)?,
             PolyOp::Pad(p) => {
                 if values.len() != 1 {
@@ -322,9 +326,8 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<
                 output_scale
             }
             PolyOp::Add => {
-                let mut scale_a = 0;
-                let scale_b = in_scales[0];
-                scale_a += in_scales[1];
+                let scale_a = in_scales[0];
+                let scale_b = in_scales[1];
                 assert_eq!(scale_a, scale_b);
                 scale_a
             }
@@ -336,19 +339,19 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<
             }
             PolyOp::Reshape(_) | PolyOp::Flatten(_) => in_scales[0],
             PolyOp::Pow(pow) => in_scales[0] * (*pow as crate::Scale),
+            PolyOp::Identity { out_scale } => out_scale.unwrap_or(in_scales[0]),
             _ => in_scales[0],
         };
         Ok(scale)
     }
 
     fn requires_homogenous_input_scales(&self) -> Vec<usize> {
-        if matches!(
-            self,
-            PolyOp::Add { .. } | PolyOp::Sub | PolyOp::Concat { .. }
-        ) {
+        if matches!(self, PolyOp::Add { .. } | PolyOp::Sub) {
             vec![0, 1]
         } else if matches!(self, PolyOp::Iff) {
             vec![1, 2]
+        } else if matches!(self, PolyOp::Concat { .. }) {
+            (0..100).collect()
         } else {
             vec![]
         }

src/circuit/ops/region.rs

@@ -1,4 +1,7 @@
-use crate::tensor::{Tensor, TensorError, TensorType, ValTensor, ValType, VarTensor};
+use crate::{
+    circuit::table::Range,
+    tensor::{Tensor, TensorError, TensorType, ValTensor, ValType, VarTensor},
+};
 use halo2_proofs::{
     circuit::Region,
     plonk::{Error, Selector},
@@ -7,9 +10,14 @@ use halo2curves::ff::PrimeField;
 use std::{
     cell::RefCell,
     collections::HashSet,
-    sync::atomic::{AtomicUsize, Ordering},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc, Mutex,
+    },
 };
 
+use super::lookup::LookupOp;
+
 /// Region error
 #[derive(Debug, thiserror::Error)]
 pub enum RegionError {
@@ -56,6 +64,8 @@ pub struct RegionCtx<'a, F: PrimeField + TensorType + PartialOrd> {
     linear_coord: usize,
     num_inner_cols: usize,
     total_constants: usize,
+    used_lookups: HashSet<LookupOp>,
+    used_range_checks: HashSet<Range>,
 }
 
 impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
@@ -75,6 +85,8 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
             row,
             linear_coord,
             total_constants: 0,
+            used_lookups: HashSet::new(),
+            used_range_checks: HashSet::new(),
         }
     }
     /// Create a new region context from a wrapped region
@@ -90,6 +102,8 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
             linear_coord,
             row,
             total_constants: 0,
+            used_lookups: HashSet::new(),
+            used_range_checks: HashSet::new(),
         }
     }
 
@@ -104,6 +118,8 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
             linear_coord,
             row,
             total_constants: 0,
+            used_lookups: HashSet::new(),
+            used_range_checks: HashSet::new(),
         }
     }
 
@@ -111,8 +127,10 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
     pub fn new_dummy_with_constants(
         row: usize,
         linear_coord: usize,
-        constants: usize,
+        total_constants: usize,
         num_inner_cols: usize,
+        used_lookups: HashSet<LookupOp>,
+        used_range_checks: HashSet<Range>,
     ) -> RegionCtx<'a, F> {
         let region = None;
         RegionCtx {
@@ -120,7 +138,9 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
             num_inner_cols,
             linear_coord,
             row,
-            total_constants: constants,
+            total_constants,
+            used_lookups,
+            used_range_checks,
         }
     }
 
@@ -170,6 +190,8 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
         let row = AtomicUsize::new(self.row());
         let linear_coord = AtomicUsize::new(self.linear_coord());
         let constants = AtomicUsize::new(self.total_constants());
+        let lookups = Arc::new(Mutex::new(self.used_lookups.clone()));
+        let range_checks = Arc::new(Mutex::new(self.used_range_checks.clone()));
 
         *output = output
             .par_enum_map(|idx, _| {
@@ -177,12 +199,16 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
                 let starting_offset = row.load(Ordering::SeqCst);
                 let starting_linear_coord = linear_coord.load(Ordering::SeqCst);
                 let starting_constants = constants.load(Ordering::SeqCst);
+                // get inner value of the locked lookups
+
                 // we need to make sure that the region is not shared between threads
                 let mut local_reg = Self::new_dummy_with_constants(
                     starting_offset,
                     starting_linear_coord,
                     starting_constants,
                     self.num_inner_cols,
+                    HashSet::new(),
+                    HashSet::new(),
                 );
                 let res = inner_loop_function(idx, &mut local_reg);
                 // we update the offset and constants
@@ -195,6 +221,11 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
                     local_reg.total_constants() - starting_constants,
                     Ordering::SeqCst,
                 );
+                // update the lookups
+                let mut lookups = lookups.lock().unwrap();
+                lookups.extend(local_reg.used_lookups());
+                let mut range_checks = range_checks.lock().unwrap();
+                range_checks.extend(local_reg.used_range_checks());
                 res
             })
             .map_err(|e| {
@@ -204,6 +235,21 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
         self.total_constants = constants.into_inner();
         self.linear_coord = linear_coord.into_inner();
         self.row = row.into_inner();
+        self.used_lookups = Arc::try_unwrap(lookups)
+            .map_err(|e| RegionError::from(format!("dummy_loop: failed to get lookups: {:?}", e)))?
+            .into_inner()
+            .map_err(|e| {
+                RegionError::from(format!("dummy_loop: failed to get lookups: {:?}", e))
+            })?;
+        self.used_range_checks = Arc::try_unwrap(range_checks)
+            .map_err(|e| {
+                RegionError::from(format!("dummy_loop: failed to get range checks: {:?}", e))
+            })?
+            .into_inner()
+            .map_err(|e| {
+                RegionError::from(format!("dummy_loop: failed to get range checks: {:?}", e))
+            })?;
+
         Ok(())
     }
 
@@ -212,15 +258,14 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
         self.region.is_none()
     }
 
-    /// duplicate_dummy
-    pub fn duplicate_dummy(&self) -> Self {
-        Self {
-            region: None,
-            linear_coord: self.linear_coord,
-            num_inner_cols: self.num_inner_cols,
-            row: self.row,
-            total_constants: self.total_constants,
-        }
+    /// add used lookup
+    pub fn add_used_lookup(&mut self, lookup: LookupOp) {
+        self.used_lookups.insert(lookup);
+    }
+
+    /// add used range check
+    pub fn add_used_range_check(&mut self, range: Range) {
+        self.used_range_checks.insert(range);
     }
 
     /// Get the offset
@@ -238,6 +283,16 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
         self.total_constants
     }
 
+    /// get used lookups
+    pub fn used_lookups(&self) -> HashSet<LookupOp> {
+        self.used_lookups.clone()
+    }
+
+    /// get used range checks
+    pub fn used_range_checks(&self) -> HashSet<Range> {
+        self.used_range_checks.clone()
+    }
+
     /// Assign a constant value
     pub fn assign_constant(&mut self, var: &VarTensor, value: F) -> Result<ValType<F>, Error> {
         self.total_constants += 1;

src/circuit/tests.rs

@@ -2154,7 +2154,7 @@ mod rangecheckpercent {
         }
 
         fn configure(cs: &mut ConstraintSystem<F>) -> Self::Config {
-            let scale = utils::F32(SCALE.pow(2) as f32);
+            let scale = utils::F32(SCALE as f32);
             let a = VarTensor::new_advice(cs, K, 1, LEN);
             let b = VarTensor::new_advice(cs, K, 1, LEN);
             let output = VarTensor::new_advice(cs, K, 1, LEN);
@@ -2162,11 +2162,12 @@ mod rangecheckpercent {
                 Self::Config::configure(cs, &[a.clone(), b.clone()], &output, CheckMode::SAFE);
             // set up a new GreaterThan and Recip tables
             let nl = &LookupOp::GreaterThan {
-                a: circuit::utils::F32((RANGE * scale.0) / 100.0),
+                a: circuit::utils::F32((RANGE * SCALE.pow(2) as f32) / 100.0),
             };
             config
                 .configure_lookup(cs, &b, &output, &a, (-32768, 32768), K, nl)
                 .unwrap();
+
             config
                 .configure_lookup(
                     cs,
@@ -2175,7 +2176,10 @@ mod rangecheckpercent {
                     &a,
                     (-32768, 32768),
                     K,
-                    &LookupOp::Recip { scale },
+                    &LookupOp::Recip {
+                        input_scale: scale,
+                        output_scale: scale,
+                    },
                 )
                 .unwrap();
             config
@@ -2511,7 +2515,8 @@ mod softmax {
                     (-32768, 32768),
                     K,
                     &LookupOp::Recip {
-                        scale: SCALE.powf(2.0).into(),
+                        input_scale: SCALE.into(),
+                        output_scale: SCALE.into(),
                     },
                 )
                 .unwrap();

src/execute.rs

@@ -630,6 +630,10 @@ pub(crate) async fn gen_witness(
     if let Some(output_path) = output {
         serde_json::to_writer(&File::create(output_path)?, &witness)?;
     }
+
+    // print the witness in debug
+    debug!("witness: \n {}", witness.as_json()?.to_colored_json_auto()?);
+
     Ok(witness)
 }
 
@@ -737,22 +741,22 @@ impl AccuracyResults {
         let median_error = errors[errors.len() / 2];
         let max_error = *errors
             .iter()
-            .max_by(|a, b| a.partial_cmp(b).unwrap())
+            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
             .unwrap();
         let min_error = *errors
             .iter()
-            .min_by(|a, b| a.partial_cmp(b).unwrap())
+            .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
             .unwrap();
 
         let mean_abs_error = abs_errors.iter().sum::<f32>() / abs_errors.len() as f32;
         let median_abs_error = abs_errors[abs_errors.len() / 2];
         let max_abs_error = *abs_errors
             .iter()
-            .max_by(|a, b| a.partial_cmp(b).unwrap())
+            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
             .unwrap();
         let min_abs_error = *abs_errors
             .iter()
-            .min_by(|a, b| a.partial_cmp(b).unwrap())
+            .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
             .unwrap();
 
         let mean_squared_error = squared_errors.iter().sum::<f32>() / squared_errors.len() as f32;

src/graph/mod.rs

@@ -968,8 +968,8 @@ impl GraphCircuit {
             lookup_safety_margin * max_lookup_inputs,
         );
         if lookup_safety_margin == 1 {
-            margin.0 += 1;
-            margin.1 += 1;
+            margin.0 += 4;
+            margin.1 += 4;
         }
 
         margin

src/graph/model.rs

@@ -80,6 +80,21 @@ pub struct ModelConfig {
 /// Representation of execution graph
 pub type NodeGraph = BTreeMap<usize, NodeType>;
 
+/// A struct for loading from an Onnx file and converting a computational graph to a circuit.
+#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
+pub struct DummyPassRes {
+    /// number of rows use
+    pub num_rows: usize,
+    /// linear coordinate
+    pub linear_coord: usize,
+    /// total const size
+    pub total_const_size: usize,
+    /// lookup ops
+    pub lookup_ops: HashSet<LookupOp>,
+    /// range checks
+    pub range_checks: HashSet<Range>,
+}
+
 /// A struct for loading from an Onnx file and converting a computational graph to a circuit.
 #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
 pub struct Model {
@@ -234,20 +249,7 @@ impl NodeType {
             NodeType::SubGraph { out_dims, .. } => out_dims.clone(),
         }
     }
-    /// Returns the lookups required by a graph
-    pub fn required_lookups(&self) -> Vec<LookupOp> {
-        match self {
-            NodeType::Node(n) => n.opkind.required_lookups(),
-            NodeType::SubGraph { model, .. } => model.required_lookups(),
-        }
-    }
-    /// Returns the lookups required by a graph
-    pub fn required_range_checks(&self) -> Vec<Range> {
-        match self {
-            NodeType::Node(n) => n.opkind.required_range_checks(),
-            NodeType::SubGraph { model, .. } => model.required_range_checks(),
-        }
-    }
+
     /// Returns the scales of the node's output.
     pub fn out_scales(&self) -> Vec<crate::Scale> {
         match self {
@@ -432,23 +434,6 @@ impl ParsedNodes {
 }
 
 impl Model {
-    fn required_lookups(&self) -> Vec<LookupOp> {
-        self.graph
-            .nodes
-            .values()
-            .flat_map(|n| n.required_lookups())
-            .collect_vec()
-    }
-
-    ///
-    fn required_range_checks(&self) -> Vec<Range> {
-        self.graph
-            .nodes
-            .values()
-            .flat_map(|n| n.required_range_checks())
-            .collect_vec()
-    }
-
     /// Creates a `Model` from a specified path to an Onnx file.
     /// # Arguments
     /// * `reader` - A reader for an Onnx file.
@@ -501,42 +486,21 @@ impl Model {
         );
         // this is the total number of variables we will need to allocate
         // for the circuit
-        let (num_rows, linear_coord, total_const_size) =
-            self.dummy_layout(run_args, &self.graph.input_shapes()?)?;
-
-        // extract the requisite lookup ops from the model
-        let mut lookup_ops: Vec<LookupOp> = self.required_lookups();
-        // extract the requisite lookup ops from the model
-        let mut range_checks: Vec<Range> = self.required_range_checks();
+        let res = self.dummy_layout(run_args, &self.graph.input_shapes()?)?;
 
         // if we're using percentage tolerance, we need to add the necessary range check ops for it.
 
-        if run_args.tolerance.val > 0.0 {
-            for scale in self.graph.get_output_scales()? {
-                let mut tolerance = run_args.tolerance;
-                tolerance.scale = scale_to_multiplier(scale).into();
-                let opkind: Box<dyn Op<Fp>> = Box::new(HybridOp::RangeCheck(tolerance));
-                lookup_ops.extend(opkind.required_lookups());
-            }
-        }
-
-        let set: HashSet<_> = lookup_ops.drain(..).collect(); // dedup
-        lookup_ops.extend(set.into_iter().sorted());
-
-        let set: HashSet<_> = range_checks.drain(..).collect(); // dedup
-        range_checks.extend(set.into_iter().sorted());
-
         Ok(GraphSettings {
             run_args: run_args.clone(),
             model_instance_shapes: instance_shapes,
             module_sizes: crate::graph::modules::ModuleSizes::default(),
-            num_rows,
-            total_assignments: linear_coord,
-            required_lookups: lookup_ops,
-            required_range_checks: range_checks,
+            num_rows: res.num_rows,
+            total_assignments: res.linear_coord,
+            required_lookups: res.lookup_ops.into_iter().collect(),
+            required_range_checks: res.range_checks.into_iter().collect(),
             model_output_scales: self.graph.get_output_scales()?,
             model_input_scales: self.graph.get_input_scales(),
-            total_const_size,
+            total_const_size: res.total_const_size,
             check_mode,
             version: env!("CARGO_PKG_VERSION").to_string(),
             num_blinding_factors: None,
@@ -591,6 +555,8 @@ impl Model {
                 inputs.iter().map(|x| x.dims()).collect::<Vec<_>>()
             );
 
+            debug!("input nodes: {:?}", n.inputs());
+
             if n.is_lookup() {
                 let (mut min, mut max) = (0, 0);
                 for i in &inputs {
@@ -1066,6 +1032,7 @@ impl Model {
                         i,
                         symbol_values,
                         run_args.div_rebasing,
+                        run_args.rebase_frac_zero_constants,
                     )?;
                     if let Some(ref scales) = override_input_scales {
                         if let Some(inp) = n.opkind.get_input() {
@@ -1523,7 +1490,7 @@ impl Model {
         &self,
         run_args: &RunArgs,
         input_shapes: &[Vec<usize>],
-    ) -> Result<(usize, usize, usize), Box<dyn Error>> {
+    ) -> Result<DummyPassRes, Box<dyn Error>> {
         info!("calculating num of constraints using dummy model layout...");
 
         let start_time = instant::Instant::now();
@@ -1608,11 +1575,15 @@ impl Model {
             region.total_constants().to_string().red()
         );
 
-        Ok((
-            region.row(),
-            region.linear_coord(),
-            region.total_constants(),
-        ))
+        let res = DummyPassRes {
+            num_rows: region.row(),
+            linear_coord: region.linear_coord(),
+            total_const_size: region.total_constants(),
+            lookup_ops: region.used_lookups(),
+            range_checks: region.used_range_checks(),
+        };
+
+        Ok(res)
     }
 
     /// Retrieves all constants from the model.

src/graph/node.rs

@@ -12,16 +12,12 @@ use crate::circuit::Constant;
 use crate::circuit::Input;
 use crate::circuit::Op;
 use crate::circuit::Unknown;
-use crate::fieldutils::felt_to_i128;
-use crate::fieldutils::i128_to_felt;
 #[cfg(not(target_arch = "wasm32"))]
 use crate::graph::new_op_from_onnx;
 use crate::tensor::Tensor;
 use crate::tensor::TensorError;
 use halo2curves::bn256::Fr as Fp;
 #[cfg(not(target_arch = "wasm32"))]
-use itertools::Itertools;
-#[cfg(not(target_arch = "wasm32"))]
 use log::trace;
 use serde::Deserialize;
 use serde::Serialize;
@@ -94,10 +90,6 @@ impl Op<Fp> for Rescaled {
         Op::<Fp>::out_scale(&*self.inner, in_scales)
     }
 
-    fn required_lookups(&self) -> Vec<LookupOp> {
-        self.inner.required_lookups()
-    }
-
     fn layout(
         &self,
         config: &mut crate::circuit::BaseConfig<Fp>,
@@ -126,14 +118,14 @@ impl Op<Fp> for Rescaled {
 pub struct RebaseScale {
     /// The operation that has to be rescaled.
     pub inner: Box<SupportedOp>,
-    /// the multiplier applied to the node output
-    pub multiplier: f64,
+    /// rebase op
+    pub rebase_op: HybridOp,
     /// scale being rebased to
     pub target_scale: i32,
     /// The original scale of the operation's inputs.
     pub original_scale: i32,
-    /// if true then the operation is a multiplicative division
-    pub div_rebasing: bool,
+    /// multiplier
+    pub multiplier: f64,
 }
 
 impl RebaseScale {
@@ -152,20 +144,27 @@ impl RebaseScale {
             let multiplier =
                 scale_to_multiplier(op_out_scale - global_scale * scale_rebase_multiplier as i32);
             if let Some(op) = inner.get_rebased() {
+                let multiplier = op.multiplier * multiplier;
                 SupportedOp::RebaseScale(RebaseScale {
                     inner: op.inner.clone(),
                     target_scale: op.target_scale,
-                    multiplier: op.multiplier * multiplier,
+                    multiplier: multiplier,
+                    rebase_op: HybridOp::Div {
+                        denom: crate::circuit::utils::F32((multiplier) as f32),
+                        use_range_check_for_int: !div_rebasing,
+                    },
                     original_scale: op.original_scale,
-                    div_rebasing,
                 })
             } else {
                 SupportedOp::RebaseScale(RebaseScale {
                     inner: Box::new(inner),
                     target_scale: global_scale * scale_rebase_multiplier as i32,
                     multiplier,
+                    rebase_op: HybridOp::Div {
+                        denom: crate::circuit::utils::F32(multiplier as f32),
+                        use_range_check_for_int: !div_rebasing,
+                    },
                     original_scale: op_out_scale,
-                    div_rebasing,
                 })
             }
         } else {
@@ -183,12 +182,16 @@ impl RebaseScale {
         if (op_out_scale < (target_scale)) && !inner.is_constant() && !inner.is_input() {
             let multiplier = scale_to_multiplier(op_out_scale - target_scale);
             if let Some(op) = inner.get_rebased() {
+                let multiplier = op.multiplier * multiplier;
                 SupportedOp::RebaseScale(RebaseScale {
                     inner: op.inner.clone(),
                     target_scale: op.target_scale,
-                    multiplier: op.multiplier * multiplier,
+                    multiplier,
                     original_scale: op.original_scale,
-                    div_rebasing,
+                    rebase_op: HybridOp::Div {
+                        denom: crate::circuit::utils::F32((multiplier) as f32),
+                        use_range_check_for_int: !div_rebasing,
+                    },
                 })
             } else {
                 SupportedOp::RebaseScale(RebaseScale {
@@ -196,22 +199,16 @@ impl RebaseScale {
                     target_scale,
                     multiplier,
                     original_scale: op_out_scale,
-                    div_rebasing,
+                    rebase_op: HybridOp::Div {
+                        denom: crate::circuit::utils::F32(multiplier as f32),
+                        use_range_check_for_int: !div_rebasing,
+                    },
                 })
             }
         } else {
             inner
         }
     }
-
-    /// Calculate the require range bracket for the operation
-    fn range_bracket(&self) -> i128 {
-        if self.div_rebasing {
-            0
-        } else {
-            self.multiplier as i128 - 1
-        }
-    }
 }
 
 impl Op<Fp> for RebaseScale {
@@ -220,28 +217,19 @@ impl Op<Fp> for RebaseScale {
     }
     fn f(&self, x: &[Tensor<Fp>]) -> Result<crate::circuit::ForwardResult<Fp>, TensorError> {
         let mut res = Op::<Fp>::f(&*self.inner, x)?;
-
-        if self.div_rebasing {
-            let ri = res.output.map(felt_to_i128);
-            let rescaled = crate::tensor::ops::nonlinearities::const_div(&ri, self.multiplier);
-            res.output = rescaled.map(i128_to_felt);
-            res.intermediate_lookups.push(ri);
-        } else {
-            let ri = res.output.map(felt_to_i128);
-            let divisor = Tensor::from(vec![self.multiplier as i128].into_iter());
-            let rescaled = crate::tensor::ops::div(&[ri, divisor.clone()])?;
-            res.output = rescaled.map(i128_to_felt);
-            res.intermediate_lookups.extend([-divisor.clone(), divisor]);
-        }
+        let rebase_res = Op::<Fp>::f(&self.rebase_op, &[res.output])?;
+        res.output = rebase_res.output;
+        res.intermediate_lookups
+            .extend(rebase_res.intermediate_lookups);
 
         Ok(res)
     }
 
     fn as_string(&self) -> String {
         format!(
-            "REBASED (div={:?}, div_r={}) ({})",
+            "REBASED (div={:?}, rebasing_op={}) ({})",
             self.multiplier,
-            self.div_rebasing,
+            <HybridOp as Op<Fp>>::as_string(&self.rebase_op),
             self.inner.as_string()
         )
     }
@@ -250,25 +238,6 @@ impl Op<Fp> for RebaseScale {
         Ok(self.target_scale)
     }
 
-    fn required_lookups(&self) -> Vec<LookupOp> {
-        let mut lookups: Vec<LookupOp> = self.inner.required_lookups();
-        if self.div_rebasing {
-            lookups.push(LookupOp::Div {
-                denom: crate::circuit::utils::F32(self.multiplier as f32),
-            });
-        }
-        lookups
-    }
-
-    fn required_range_checks(&self) -> Vec<crate::circuit::table::Range> {
-        let mut range_checks = self.inner.required_range_checks();
-        if !self.div_rebasing {
-            let bracket = self.range_bracket();
-            range_checks.push((-bracket, bracket));
-        }
-        range_checks
-    }
-
     fn layout(
         &self,
         config: &mut crate::circuit::BaseConfig<Fp>,
@@ -278,25 +247,8 @@ impl Op<Fp> for RebaseScale {
         let original_res = self
             .inner
             .layout(config, region, values)?
-            .ok_or("no layout")?;
-
-        if !self.div_rebasing {
-            Ok(Some(crate::circuit::layouts::div(
-                config,
-                region,
-                &[original_res],
-                Fp::from(self.multiplier as u64),
-            )?))
-        } else {
-            Ok(Some(crate::circuit::layouts::nonlinearity(
-                config,
-                region,
-                &[original_res],
-                &LookupOp::Div {
-                    denom: crate::circuit::utils::F32(self.multiplier as f32),
-                },
-            )?))
-        }
+            .ok_or("no inner layout")?;
+        self.rebase_op.layout(config, region, &[original_res])
     }
 
     fn clone_dyn(&self) -> Box<dyn Op<Fp>> {
@@ -479,14 +431,6 @@ impl Op<Fp> for SupportedOp {
         self
     }
 
-    fn required_lookups(&self) -> Vec<LookupOp> {
-        self.as_op().required_lookups()
-    }
-
-    fn required_range_checks(&self) -> Vec<crate::circuit::table::Range> {
-        self.as_op().required_range_checks()
-    }
-
     fn out_scale(&self, in_scales: Vec<crate::Scale>) -> Result<crate::Scale, Box<dyn Error>> {
         self.as_op().out_scale(in_scales)
     }
@@ -520,15 +464,7 @@ impl Tabled for Node {
 
     fn headers() -> Vec<std::borrow::Cow<'static, str>> {
         let mut headers = Vec::with_capacity(Self::LENGTH);
-        for i in [
-            "idx",
-            "opkind",
-            "out_scale",
-            "inputs",
-            "out_dims",
-            "required_lookups",
-            "required_range_checks",
-        ] {
+        for i in ["idx", "opkind", "out_scale", "inputs", "out_dims"] {
             headers.push(std::borrow::Cow::Borrowed(i));
         }
         headers
@@ -541,18 +477,6 @@ impl Tabled for Node {
         fields.push(std::borrow::Cow::Owned(self.out_scale.to_string()));
         fields.push(std::borrow::Cow::Owned(display_vector(&self.inputs)));
         fields.push(std::borrow::Cow::Owned(display_vector(&self.out_dims)));
-        fields.push(std::borrow::Cow::Owned(format!(
-            "{:?}",
-            self.opkind
-                .required_lookups()
-                .iter()
-                .map(<LookupOp as Op<Fp>>::as_string)
-                .collect_vec()
-        )));
-        fields.push(std::borrow::Cow::Owned(format!(
-            "{:?}",
-            self.opkind.required_range_checks()
-        )));
         fields
     }
 }
@@ -583,9 +507,8 @@ impl Node {
         idx: usize,
         symbol_values: &SymbolValues,
         div_rebasing: bool,
+        rebase_frac_zero_constants: bool,
     ) -> Result<Self, Box<dyn Error>> {
-        use log::warn;
-
         trace!("Create {:?}", node);
         trace!("Create op {:?}", node.op);
 
@@ -623,6 +546,7 @@ impl Node {
             node.clone(),
             &mut inputs,
             symbol_values,
+            rebase_frac_zero_constants,
         )?; // parses the op name
 
         // we can only take the inputs as mutable once -- so we need to collect them first
@@ -678,8 +602,6 @@ impl Node {
                     input_node.bump_scale(out_scale);
                     in_scales[input] = out_scale;
                 }
-            } else {
-                warn!("input {} not found for rescaling, skipping ...", input);
             }
         }
 

src/graph/utilities.rs

@@ -243,6 +243,7 @@ pub fn new_op_from_onnx(
     node: OnnxNode<TypedFact, Box<dyn TypedOp>>,
     inputs: &mut [super::NodeType],
     symbol_values: &SymbolValues,
+    rebase_frac_zero_constants: bool,
 ) -> Result<(SupportedOp, Vec<usize>), Box<dyn std::error::Error>> {
     use crate::circuit::InputType;
 
@@ -261,7 +262,9 @@ pub fn new_op_from_onnx(
             inputs[index].bump_scale(scale);
             c.rebase_scale(scale)?;
             inputs[index].replace_opkind(SupportedOp::Constant(c.clone()));
-            Ok(SupportedOp::Linear(PolyOp::Identity))
+            Ok(SupportedOp::Linear(PolyOp::Identity {
+                out_scale: Some(scale),
+            }))
         } else {
             Ok(default_op)
         }
@@ -282,8 +285,8 @@ pub fn new_op_from_onnx(
                         "shift left".to_string(),
                     )));
                 }
-                SupportedOp::Nonlinear(LookupOp::Div {
-                    denom: crate::circuit::utils::F32(1.0 / 2.0f32.powf(raw_values[0])),
+                SupportedOp::Linear(PolyOp::Identity {
+                    out_scale: Some(input_scales[0] - raw_values[0] as i32),
                 })
             } else {
                 return Err(Box::new(GraphError::OpMismatch(
@@ -304,8 +307,8 @@ pub fn new_op_from_onnx(
                         "shift right".to_string(),
                     )));
                 }
-                SupportedOp::Nonlinear(LookupOp::Div {
-                    denom: crate::circuit::utils::F32(2.0f32.powf(raw_values[0])),
+                SupportedOp::Linear(PolyOp::Identity {
+                    out_scale: Some(input_scales[0] + raw_values[0] as i32),
                 })
             } else {
                 return Err(Box::new(GraphError::OpMismatch(
@@ -544,7 +547,7 @@ pub fn new_op_from_onnx(
             // Raw values are always f32
             let raw_value = extract_tensor_value(op.0)?;
             // If bool or a tensor dimension then don't scale
-            let constant_scale = match dt {
+            let mut constant_scale = match dt {
                 DatumType::Bool
                 | DatumType::TDim
                 | DatumType::I64
@@ -559,6 +562,12 @@ pub fn new_op_from_onnx(
                 _ => return Err(Box::new(GraphError::UnsupportedDataType)),
             };
 
+            // if all raw_values are round then set scale to 0
+            let all_round = raw_value.iter().all(|x| (x).fract() == 0.0);
+            if all_round && rebase_frac_zero_constants {
+                constant_scale = 0;
+            }
+
             // Quantize the raw value
             let quantized_value =
                 quantize_tensor(raw_value.clone(), constant_scale, param_visibility)?;
@@ -665,8 +674,10 @@ pub fn new_op_from_onnx(
                 if unit == 0. {
                     SupportedOp::Nonlinear(LookupOp::ReLU)
                 } else {
+                    // get the non-constant index
+                    let non_const_idx = if const_idx == 0 { 1 } else { 0 };
                     SupportedOp::Nonlinear(LookupOp::Max {
-                        scale: scale_to_multiplier(inputs[0].out_scales()[0]).into(),
+                        scale: scale_to_multiplier(inputs[non_const_idx].out_scales()[0]).into(),
                         a: crate::circuit::utils::F32(unit),
                     })
                 }
@@ -707,8 +718,11 @@ pub fn new_op_from_onnx(
                     deleted_indices.push(const_idx);
                 }
 
+                // get the non-constant index
+                let non_const_idx = if const_idx == 0 { 1 } else { 0 };
+
                 SupportedOp::Nonlinear(LookupOp::Min {
-                    scale: scale_to_multiplier(inputs[0].out_scales()[0]).into(),
+                    scale: scale_to_multiplier(inputs[non_const_idx].out_scales()[0]).into(),
                     a: crate::circuit::utils::F32(unit),
                 })
             } else {
@@ -717,16 +731,12 @@ pub fn new_op_from_onnx(
         }
         "Recip" => {
             let in_scale = inputs[0].out_scales()[0];
+            let max_scale = std::cmp::max(scales.get_max(), in_scale);
             // If the input scale is larger than the params scale
-            let scale_diff = std::cmp::max(scales.input, scales.params) - inputs[0].out_scales()[0];
-            let additional_scale = if scale_diff > 0 {
-                scale_to_multiplier(scale_diff)
-            } else {
-                1.0
-            };
-
-            SupportedOp::Nonlinear(LookupOp::Recip {
-                scale: (scale_to_multiplier(in_scale).powf(2.0) * additional_scale).into(),
+            SupportedOp::Hybrid(HybridOp::Recip {
+                input_scale: (scale_to_multiplier(in_scale) as f32).into(),
+                output_scale: (scale_to_multiplier(max_scale) as f32).into(),
+                use_range_check_for_int: false,
             })
         }
 
@@ -751,7 +761,9 @@ pub fn new_op_from_onnx(
         "Scan" => {
             return Err("scan should never be analyzed explicitly".into());
         }
-        "QuantizeLinearU8" | "DequantizeLinearF32" => SupportedOp::Linear(PolyOp::Identity),
+        "QuantizeLinearU8" | "DequantizeLinearF32" => {
+            SupportedOp::Linear(PolyOp::Identity { out_scale: None })
+        }
         "Abs" => SupportedOp::Nonlinear(LookupOp::Abs),
         "Neg" => SupportedOp::Linear(PolyOp::Neg),
         "Sigmoid" => SupportedOp::Nonlinear(LookupOp::Sigmoid {
@@ -856,11 +868,11 @@ pub fn new_op_from_onnx(
                             }),
                         )?
                     } else {
-                        SupportedOp::Linear(PolyOp::Identity)
+                        SupportedOp::Linear(PolyOp::Identity { out_scale: None })
                     }
                 }
                 DatumType::F16 | DatumType::F32 | DatumType::F64 => {
-                    SupportedOp::Linear(PolyOp::Identity)
+                    SupportedOp::Linear(PolyOp::Identity { out_scale: None })
                 }
                 _ => return Err(Box::new(GraphError::UnsupportedDataType)),
             }
@@ -885,12 +897,15 @@ pub fn new_op_from_onnx(
                 let const_idx = const_idx[0];
                 if let Some(c) = inputs[const_idx].opkind().get_mutable_constant() {
                     if c.raw_values.len() == 1 && c.raw_values[0] < 1. {
-                        inputs[const_idx].decrement_use();
-                        deleted_indices.push(const_idx);
-                        op = SupportedOp::Nonlinear(LookupOp::Div {
-                            // we invert the constant for division
-                            denom: crate::circuit::utils::F32(1. / c.raw_values[0]),
-                        })
+                        // if not divisible by 2 then we need to add a range check
+                        let raw_values = 1.0 / c.raw_values[0];
+                        if raw_values.log2().fract() == 0.0 {
+                            inputs[const_idx].decrement_use();
+                            deleted_indices.push(const_idx);
+                            op = SupportedOp::Linear(PolyOp::Identity {
+                                out_scale: Some(input_scales[0] + raw_values.log2() as i32),
+                            });
+                        }
                     }
                 }
             }

src/graph/vars.rs

@@ -237,6 +237,11 @@ impl VarScales {
         std::cmp::max(self.input, self.params)
     }
 
+    ///
+    pub fn get_min(&self) -> crate::Scale {
+        std::cmp::min(self.input, self.params)
+    }
+
     /// Place in [VarScales] struct.
     pub fn from_args(args: &RunArgs) -> Result<Self, Box<dyn Error>> {
         Ok(Self {

src/lib.rs

@@ -111,8 +111,11 @@ pub struct RunArgs {
     #[arg(long, default_value = "private")]
     pub param_visibility: Visibility,
     #[arg(long, default_value = "false")]
-    /// Multiplicative division
+    /// Rebase the scale using lookup table for division instead of using a range check
     pub div_rebasing: bool,
+    /// Should constants with 0.0 fraction be rebased to scale 0
+    #[arg(long, default_value = "false")]
+    pub rebase_frac_zero_constants: bool,
     /// check mode (safe, unsafe, etc)
     #[arg(long, default_value = "unsafe")]
     pub check_mode: CheckMode,
@@ -133,6 +136,7 @@ impl Default for RunArgs {
             output_visibility: Visibility::Public,
             param_visibility: Visibility::Private,
             div_rebasing: false,
+            rebase_frac_zero_constants: false,
             check_mode: CheckMode::UNSAFE,
         }
     }

src/python.rs

@@ -162,6 +162,8 @@ struct PyRunArgs {
     #[pyo3(get, set)]
     pub div_rebasing: bool,
     #[pyo3(get, set)]
+    pub rebase_frac_zero_constants: bool,
+    #[pyo3(get, set)]
     pub check_mode: CheckMode,
 }
 
@@ -190,6 +192,7 @@ impl From<PyRunArgs> for RunArgs {
             param_visibility: py_run_args.param_visibility,
             variables: py_run_args.variables,
             div_rebasing: py_run_args.div_rebasing,
+            rebase_frac_zero_constants: py_run_args.rebase_frac_zero_constants,
             check_mode: py_run_args.check_mode,
         }
     }
@@ -210,6 +213,7 @@ impl Into<PyRunArgs> for RunArgs {
             param_visibility: self.param_visibility,
             variables: self.variables,
             div_rebasing: self.div_rebasing,
+            rebase_frac_zero_constants: self.rebase_frac_zero_constants,
             check_mode: self.check_mode,
         }
     }

src/tensor/ops.rs

@@ -992,45 +992,6 @@ pub fn mult<T: TensorType + Mul<Output = T> + std::marker::Send + std::marker::S
     Ok(output)
 }
 
-/// Divides multiple tensors.
-/// # Arguments
-/// * `t` - Tensors
-/// # Examples
-/// ```
-/// use ezkl::tensor::Tensor;
-/// use ezkl::tensor::ops::div;
-/// let x = Tensor::<i128>::new(
-///    Some(&[2, 1, 2, 1, 1, 1]),
-/// &[2, 3],
-/// ).unwrap();
-/// let k = Tensor::<i128>::new(
-///   Some(&[2, 3, 2, 1, 1, 1]),
-/// &[2, 3],
-/// ).unwrap();
-/// let result = div(&[x, k]).unwrap();
-/// let expected = Tensor::<i128>::new(Some(&[1, 0, 1, 1, 1, 1]), &[2, 3]).unwrap();
-/// assert_eq!(result, expected);
-/// ```
-pub fn div<
-    T: TensorType
-        + Div<Output = T>
-        + Mul<Output = T>
-        + From<u64>
-        + std::marker::Send
-        + std::marker::Sync,
->(
-    t: &[Tensor<T>],
-) -> Result<Tensor<T>, TensorError> {
-    // calculate value of output
-    let mut output: Tensor<T> = t[0].clone();
-
-    for e in t[1..].iter() {
-        output = (output / e.clone())?;
-    }
-
-    Ok(output)
-}
-
 /// Rescale a tensor with a const integer (similar to const_mult).
 /// # Arguments
 ///
@@ -3164,7 +3125,7 @@ pub mod nonlinearities {
 
         let sum = sum(&exp).unwrap();
         intermediate_values.push(sum.clone());
-        let inv_denom = recip(&sum, scale.powf(2.0));
+        let inv_denom = recip(&sum, scale, scale);
 
         ((exp * inv_denom).unwrap(), intermediate_values)
     }
@@ -3201,7 +3162,7 @@ pub mod nonlinearities {
         // the more accurate calculation is commented out and we implement as below so it matches the steps in layout
         let scale = input_scale * output_scale;
         let diff: Tensor<i128> = sub(t).unwrap();
-        let recip = recip(&t[0], scale as f64);
+        let recip = recip(&t[0], input_scale as f64, output_scale as f64);
         let product = mult(&[diff, recip]).unwrap();
         let _tol = ((tol / 100.0) * scale as f32).round() as f64;
         let upper_bound = greater_than(&product, _tol);
@@ -3812,14 +3773,15 @@ pub mod nonlinearities {
     ///     &[2, 3],
     /// ).unwrap();
     /// let k = 2_f64;
-    /// let result = recip(&x, k);
+    /// let result = recip(&x, 1.0, k);
     /// let expected = Tensor::<i128>::new(Some(&[1, 2, 1, 0, 2, 2]), &[2, 3]).unwrap();
     /// assert_eq!(result, expected);
     /// ```
-    pub fn recip(a: &Tensor<i128>, scale: f64) -> Tensor<i128> {
+    pub fn recip(a: &Tensor<i128>, input_scale: f64, out_scale: f64) -> Tensor<i128> {
         a.par_enum_map(|_, a_i| {
-            let denom = (1_f64) / (a_i as f64 + f64::EPSILON);
-            let d_inv_x = scale * denom;
+            let rescaled = (a_i as f64) / input_scale;
+            let denom = (1_f64) / (rescaled + f64::EPSILON);
+            let d_inv_x = out_scale * denom;
             Ok::<_, TensorError>(d_inv_x.round() as i128)
         })
         .unwrap()

tests/integration_tests.rs

@@ -182,12 +182,13 @@ mod native_tests {
         "mnist_gan",
     ];
 
-    const ACCURACY_CAL_TESTS: [&str; 5] = [
+    const ACCURACY_CAL_TESTS: [&str; 6] = [
         "accuracy",
         "1l_mlp",
         "4l_relu_conv_fc",
         "1l_elu",
         "1l_prelu",
+        "1l_tiny_div",
     ];
 
     const TESTS: [&str; 77] = [
@@ -489,7 +490,7 @@ mod native_tests {
                 test_dir.close().unwrap();
             }
 
-            seq!(N in 0..=4 {
+            seq!(N in 0..=5 {
             #(#[test_case(ACCURACY_CAL_TESTS[N])])*
             fn mock_accuracy_cal_tests(test: &str) {
                 crate::native_tests::init_binary();
@@ -2032,7 +2033,7 @@ mod native_tests {
             1,
             "resources",
             // we need the accuracy
-            Some(vec![7, 8]),
+            Some(vec![4]),
             1,
             false,
         );

tests/output_comparison.py

@@ -78,14 +78,20 @@ def compare_outputs(zk_output, onnx_output):
 
     zip_object = zip(np.array(zk_output).flatten(),
                      np.array(onnx_output).flatten())
-    for list1_i, list2_i in zip_object:
+    for (i, (list1_i, list2_i)) in enumerate(zip_object):
         if list1_i == 0.0 and list2_i == 0.0:
             res.append(0)
         else:
             diff = list1_i - list2_i
             res.append(100 * (diff) / (list2_i))
+            # iterate and print the diffs  if they are greater than 0.0
+            if abs(diff) > 0.0:
+                print("------- index: ", i)
+                print("------- diff: ", diff)
+                print("------- zk_output: ", list1_i)
+                print("------- onnx_output: ", list2_i)
+
 
-    print("res: ", res)
 
     return np.mean(np.abs(res))
 

tests/wasm/settings.json

@@ -23,6 +23,7 @@
         "output_visibility": "Public",
         "param_visibility": "Private",
         "div_rebasing": false,
+        "rebase_frac_zero_constants": false,
         "check_mode": "UNSAFE"
     },
     "num_rows": 16,