feat: lookupless sqrt and rsqrt

.github/workflows/rust.yml

@@ -240,6 +240,8 @@ jobs:
           locked: true
       # - name: The Worm Mock
       #   run: cargo nextest run --release --verbose tests::large_mock_::large_tests_5_expects -- --include-ignored
+      - name: public outputs and bounded lookup log
+        run: cargo nextest run --release --verbose tests::mock_bounded_lookup_log --test-threads 32
       - name: public outputs and tolerance > 0
         run: cargo nextest run --release --verbose tests::mock_tolerance_public_outputs_ --test-threads 32
       - name: public outputs + batch size == 10

examples/notebooks/data_attest_hashed.ipynb

@@ -648,7 +648,7 @@
             "name": "python",
             "nbconvert_exporter": "python",
             "pygments_lexer": "ipython3",
-            "version": "3.12.5"
+            "version": "3.12.7"
         },
         "orig_nbformat": 4
     },

examples/notebooks/simple_demo_aggregated_proofs.ipynb

@@ -171,7 +171,7 @@
                 "json.dump(data, open(cal_path, 'w'))\n",
                 "\n",
                 "\n",
-                "ezkl.calibrate_settings(cal_path, model_path, settings_path, \"resources\")"
+                "await ezkl.calibrate_settings(cal_path, model_path, settings_path, \"resources\")"
             ]
         },
         {
@@ -328,7 +328,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 26,
             "id": "171702d3",
             "metadata": {},
             "outputs": [],
@@ -348,7 +348,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 27,
             "id": "671dfdd5",
             "metadata": {},
             "outputs": [],
@@ -364,7 +364,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 28,
             "id": "50eba2f4",
             "metadata": {},
             "outputs": [],
@@ -399,9 +399,9 @@
             "name": "python",
             "nbconvert_exporter": "python",
             "pygments_lexer": "ipython3",
-            "version": "3.9.15"
+            "version": "3.12.7"
         }
     },
     "nbformat": 4,
     "nbformat_minor": 5
-}
+}

src/bindings/python.rs

@@ -180,9 +180,6 @@ struct PyRunArgs {
     /// list[tuple[str, int]]: Hand-written parser for graph variables, eg. batch_size=1
     pub variables: Vec<(String, usize)>,
     #[pyo3(get, set)]
-    /// bool: Rebase the scale using lookup table for division instead of using a range check
-    pub div_rebasing: bool,
-    #[pyo3(get, set)]
     /// bool: Should constants with 0.0 fraction be rebased to scale 0
     pub rebase_frac_zero_constants: bool,
     #[pyo3(get, set)]
@@ -227,7 +224,6 @@ impl From<PyRunArgs> for RunArgs {
             output_visibility: py_run_args.output_visibility,
             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,
             commitment: Some(py_run_args.commitment.into()),
@@ -252,7 +248,6 @@ impl Into<PyRunArgs> for RunArgs {
             output_visibility: self.output_visibility,
             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,
             commitment: self.commitment.into(),
@@ -878,8 +873,6 @@ fn gen_settings(
 /// max_logrows: int
 ///     Optional max logrows to use for calibration
 ///
-/// only_range_check_rebase: bool
-///     Check ranges when rebasing
 ///
 /// Returns
 /// -------
@@ -894,7 +887,6 @@ fn gen_settings(
     scales = None,
     scale_rebase_multiplier = DEFAULT_SCALE_REBASE_MULTIPLIERS.split(",").map(|x| x.parse().unwrap()).collect(),
     max_logrows = None,
-    only_range_check_rebase = DEFAULT_ONLY_RANGE_CHECK_REBASE.parse().unwrap(),
 ))]
 fn calibrate_settings(
     py: Python,
@@ -906,7 +898,6 @@ fn calibrate_settings(
     scales: Option<Vec<crate::Scale>>,
     scale_rebase_multiplier: Vec<u32>,
     max_logrows: Option<u32>,
-    only_range_check_rebase: bool,
 ) -> PyResult<Bound<'_, PyAny>> {
     pyo3_asyncio::tokio::future_into_py(py, async move {
         crate::execute::calibrate(
@@ -917,7 +908,6 @@ fn calibrate_settings(
             lookup_safety_margin,
             scales,
             scale_rebase_multiplier,
-            only_range_check_rebase,
             max_logrows,
         )
         .await

src/circuit/ops/hybrid.rs

@@ -16,7 +16,13 @@ pub enum HybridOp {
     Ln {
         scale: utils::F32,
     },
-
+    Rsqrt {
+        input_scale: utils::F32,
+        output_scale: utils::F32,
+    },
+    Sqrt {
+        scale: utils::F32,
+    },
     RoundHalfToEven {
         scale: utils::F32,
         legs: usize,
@@ -39,7 +45,6 @@ pub enum HybridOp {
     },
     Div {
         denom: utils::F32,
-        use_range_check_for_int: bool,
     },
     ReduceMax {
         axes: Vec<usize>,
@@ -116,6 +121,14 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Hybrid
 
     fn as_string(&self) -> String {
         match self {
+            HybridOp::Rsqrt {
+                input_scale,
+                output_scale,
+            } => format!(
+                "RSQRT (input_scale={}, output_scale={})",
+                input_scale, output_scale
+            ),
+            HybridOp::Sqrt { scale } => format!("SQRT(scale={})", scale),
             HybridOp::Ln { scale } => format!("LN(scale={})", scale),
             HybridOp::RoundHalfToEven { scale, legs } => {
                 format!("ROUND_HALF_TO_EVEN(scale={}, legs={})", scale, legs)
@@ -124,22 +137,16 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Hybrid
             HybridOp::Floor { scale, legs } => format!("FLOOR(scale={}, legs={})", scale, legs),
             HybridOp::Round { scale, legs } => format!("ROUND(scale={}, legs={})", scale, legs),
 
-            HybridOp::Max => format!("MAX"),
-            HybridOp::Min => format!("MIN"),
+            HybridOp::Max => "MAX".to_string(),
+            HybridOp::Min => "MIN".to_string(),
             HybridOp::Recip {
                 input_scale,
                 output_scale,
             } => format!(
                 "RECIP (input_scale={}, output_scale={})",
                 input_scale, output_scale
             ),
-            HybridOp::Div {
-                denom,
-                use_range_check_for_int,
-            } => format!(
-                "DIV (denom={}, use_range_check_for_int={})",
-                denom, use_range_check_for_int
-            ),
+            HybridOp::Div { denom } => format!("DIV (denom={})", denom),
             HybridOp::SumPool {
                 padding,
                 stride,
@@ -194,6 +201,19 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Hybrid
         values: &[ValTensor<F>],
     ) -> Result<Option<ValTensor<F>>, CircuitError> {
         Ok(Some(match self {
+            HybridOp::Rsqrt {
+                input_scale,
+                output_scale,
+            } => layouts::rsqrt(
+                config,
+                region,
+                values[..].try_into()?,
+                *input_scale,
+                *output_scale,
+            )?,
+            HybridOp::Sqrt { scale } => {
+                layouts::sqrt(config, region, values[..].try_into()?, *scale)?
+            }
             HybridOp::Ln { scale } => layouts::ln(config, region, values[..].try_into()?, *scale)?,
             HybridOp::RoundHalfToEven { scale, legs } => {
                 layouts::round_half_to_even(config, region, values[..].try_into()?, *scale, *legs)?
@@ -233,13 +253,9 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Hybrid
                 integer_rep_to_felt(input_scale.0 as i128),
                 integer_rep_to_felt(output_scale.0 as i128),
             )?,
-            HybridOp::Div {
-                denom,
-                use_range_check_for_int,
-                ..
-            } => {
-                if denom.0.fract() == 0.0 && *use_range_check_for_int {
-                    layouts::loop_div(
+            HybridOp::Div { denom, .. } => {
+                if denom.0.fract() == 0.0 {
+                    layouts::div(
                         config,
                         region,
                         values[..].try_into()?,
@@ -330,9 +346,15 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Hybrid
             | HybridOp::ReduceArgMax { .. }
             | HybridOp::OneHot { .. }
             | HybridOp::ReduceArgMin { .. } => 0,
-            HybridOp::Softmax { output_scale, .. } | HybridOp::Recip { output_scale, .. } => {
+
+            HybridOp::Recip { output_scale, .. } | HybridOp::Rsqrt { output_scale, .. } => {
                 multiplier_to_scale(output_scale.0 as f64)
             }
+            HybridOp::Softmax {
+                output_scale,
+                input_scale,
+                ..
+            } => multiplier_to_scale((output_scale.0 * input_scale.0) as f64),
             HybridOp::Ln {
                 scale: output_scale,
             } => 4 * multiplier_to_scale(output_scale.0 as f64),