fix: calibrate from total min/max on lookups rather than individual x (#690)

.github/workflows/rust.yml

@@ -465,7 +465,7 @@ jobs:
       #   run: cargo nextest run --release --verbose tests::kzg_fuzz_ --test-threads 6
 
   prove-and-verify-mock-aggr-tests:
-    runs-on: ubuntu-latest-32-cores
+    runs-on: self-hosted
     needs: [build, library-tests]
     steps:
       - uses: actions/checkout@v4
@@ -611,27 +611,7 @@ jobs:
         run: source .env/bin/activate; cargo nextest run --release --verbose tests::resources_accuracy_measurement_public_outputs_
 
   python-integration-tests:
-    runs-on:
-      large-self-hosted
-      # Service containers to run with `container-job`
-    services:
-      # Label used to access the service container
-      postgres:
-        # Docker Hub image
-        image: postgres
-        env:
-          POSTGRES_USER: ubuntu
-          POSTGRES_HOST_AUTH_METHOD: trust
-        # Set health checks to wait until postgres has started
-        options: >-
-          --health-cmd pg_isready
-          --health-interval 10s
-          --health-timeout 5s
-          --health-retries 5
-        ports:
-          # Maps tcp port 5432 on service container to the host
-          - 5432:5432
-    # needs: [build,  library-tests, docs]
+    runs-on: large-self-hosted
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v4
@@ -663,13 +643,14 @@ jobs:
       #     # now dump the contents of the file into a file called kaggle.json
       #     echo $KAGGLE_API_KEY > /home/ubuntu/.kaggle/kaggle.json
       #     chmod 600 /home/ubuntu/.kaggle/kaggle.json
-      - name: Tictactoe tutorials
-        run: source .env/bin/activate; cargo nextest run py_tests::tests::tictactoe_ --test-threads 1
-      # - name: Postgres tutorials
-      #   run: source .env/bin/activate; cargo nextest run py_tests::tests::postgres_ --test-threads 1
-      - name: All notebooks
-        run: source .env/bin/activate; cargo nextest run py_tests::tests::run_notebook_ --test-threads 1
       - name: NBEATS tutorial
         run: source .env/bin/activate; cargo nextest run py_tests::tests::nbeats_
       - name: Voice tutorial
         run: source .env/bin/activate; cargo nextest run py_tests::tests::voice_
+      - name: All notebooks
+        run: source .env/bin/activate; cargo nextest run py_tests::tests::run_notebook_ --test-threads 1
+      - name: Tictactoe tutorials
+        run: source .env/bin/activate; cargo nextest run py_tests::tests::tictactoe_ --test-threads 1
+      # - name: Postgres tutorials
+      #   run: source .env/bin/activate; cargo nextest run py_tests::tests::postgres_ --test-threads 1
+     

examples/notebooks/ezkl_demo.ipynb

@@ -42,7 +42,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": null,
       "metadata": {
         "id": "gvQ5HL1bTDWF"
       },
@@ -441,7 +441,9 @@
         "# Serialize calibration data into file:\n",
         "json.dump(data, open(cal_data_path, 'w'))\n",
         "\n",
-        "res = ezkl.calibrate_settings(cal_data_path, model_path, settings_path, \"resources\")  # Optimize for resources"
+        "# Optimize for resources, we cap logrows at 17 to reduce setup and proving time, at the expense of accuracy\n",
+        "# You may want to increase the max logrows if accuracy is a concern\n",
+        "res = ezkl.calibrate_settings(cal_data_path, model_path, settings_path, \"resources\", max_logrows = 12, scales = [2])"
       ]
     },
     {
@@ -506,9 +508,8 @@
         "        compiled_model_path,\n",
         "        vk_path,\n",
         "        pk_path,\n",
-        "        \n",
         "    )\n",
-        " \n",
+        "\n",
         "\n",
         "assert res == True\n",
         "assert os.path.isfile(vk_path)\n",
@@ -563,7 +564,6 @@
         "        compiled_model_path,\n",
         "        pk_path,\n",
         "        proof_path,\n",
-        "        \n",
         "        \"single\",\n",
         "    )\n",
         "\n",

examples/notebooks/nbeats_timeseries_forecasting.ipynb

@@ -780,7 +780,7 @@
         "pk_path = os.path.join('test.pk')\n",
         "vk_path = os.path.join('test.vk')\n",
         "settings_path = os.path.join('settings.json')\n",
-        "",
+        "\n",
         "witness_path = os.path.join('witness.json')\n",
         "data_path = os.path.join('input.json')"
       ]
@@ -845,7 +845,7 @@
         "res = ezkl.gen_settings(model_path, settings_path)\n",
         "assert res == True\n",
         "\n",
-        "res = ezkl.calibrate_settings(data_path, model_path, settings_path, \"resources\", max_logrows = 20, scales = [5,6])\n",
+        "res = ezkl.calibrate_settings(data_path, model_path, settings_path, \"resources\", max_logrows = 20, scales = [3])\n",
         "assert res == True"
       ]
     },
@@ -887,11 +887,28 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 30,
       "metadata": {
         "id": "12YIcFr85X9-"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "spawning module 2\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "quotient_poly_degree 4\n",
+            "n 262144\n",
+            "extended_k 20\n"
+          ]
+        }
+      ],
       "source": [
         "res = ezkl.setup(\n",
         "        compiled_model_path,\n",
@@ -971,9 +988,9 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.9.13"
+      "version": "3.9.15"
     }
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
+}

examples/notebooks/variance.ipynb

@@ -237,7 +237,7 @@
         "\n",
         "ezkl.gen_settings(onnx_filename, settings_filename)\n",
         "ezkl.calibrate_settings(\n",
-        "    input_filename, onnx_filename, settings_filename, \"resources\")\n",
+        "    input_filename, onnx_filename, settings_filename, \"resources\", scales = [4])\n",
         "res = ezkl.get_srs(settings_filename)\n",
         "ezkl.compile_circuit(onnx_filename, compiled_filename, settings_filename)\n",
         "\n",
@@ -255,7 +255,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 7,
       "metadata": {
         "id": "fULvvnK7_CMb"
       },
@@ -451,7 +451,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.9.13"
+      "version": "3.9.15"
     }
   },
   "nbformat": 4,

src/execute.rs

@@ -823,8 +823,6 @@ pub(crate) fn calibrate(
             "input scale: {}, param scale: {}, scale rebase multiplier: {}",
             input_scale, param_scale, scale_rebase_multiplier
         ));
-        // vec of settings copied chunks.len() times
-        let run_args_iterable = vec![settings.run_args.clone(); chunks.len()];
 
         #[cfg(unix)]
         let _r = match Gag::stdout() {
@@ -840,37 +838,30 @@ pub(crate) fn calibrate(
         let key = (input_scale, param_scale, scale_rebase_multiplier);
         forward_pass_res.insert(key, vec![]);
 
-        let tasks = chunks
+        let local_run_args = RunArgs {
+            input_scale,
+            param_scale,
+            scale_rebase_multiplier,
+            ..settings.run_args.clone()
+        };
+
+        let mut circuit = match GraphCircuit::from_run_args(&local_run_args, &model_path) {
+            Ok(c) => c,
+            Err(_) => return Err(format!("failed to create circuit from run args").into()),
+        };
+
+        chunks
             .iter()
-            .zip(run_args_iterable)
-            .map(|(chunk, run_args)| {
-                // we need to create a new run args for each chunk
-                // time it
+            .map(|chunk| {
                 let chunk = chunk.clone();
-                let local_run_args = RunArgs {
-                    input_scale,
-                    param_scale,
-                    scale_rebase_multiplier,
-                    ..run_args.clone()
-                };
-
-                let original_settings = settings.clone();
-
-                let mut circuit = match GraphCircuit::from_run_args(&local_run_args, &model_path) {
-                    Ok(c) => c,
-                    Err(_) => {
-                        return Err(format!("failed to create circuit from run args"))
-                            as Result<GraphSettings, String>
-                    }
-                };
 
                 let data = circuit
                     .load_graph_from_file_exclusively(&chunk)
                     .map_err(|e| format!("failed to load circuit inputs: {}", e))?;
 
                 let forward_res = circuit
-                    .calibrate(&data, max_logrows, lookup_safety_margin)
-                    .map_err(|e| format!("failed to calibrate: {}", e))?;
+                    .forward(&mut data.clone(), None, None)
+                    .map_err(|e| format!("failed to forward: {}", e))?;
 
                 // push result to the hashmap
                 forward_pass_res
@@ -878,38 +869,30 @@ pub(crate) fn calibrate(
                     .ok_or("key not found")?
                     .push(forward_res);
 
-                let settings = circuit.settings().clone();
-
-                let found_run_args = RunArgs {
-                    input_scale: settings.run_args.input_scale,
-                    param_scale: settings.run_args.param_scale,
-                    lookup_range: settings.run_args.lookup_range,
-                    logrows: settings.run_args.logrows,
-                    scale_rebase_multiplier: settings.run_args.scale_rebase_multiplier,
-                    ..run_args.clone()
-                };
-
-                let found_settings = GraphSettings {
-                    run_args: found_run_args,
-                    required_lookups: settings.required_lookups,
-                    model_output_scales: settings.model_output_scales,
-                    model_input_scales: settings.model_input_scales,
-                    num_rows: settings.num_rows,
-                    total_assignments: settings.total_assignments,
-                    total_const_size: settings.total_const_size,
-                    ..original_settings.clone()
-                };
-
-                Ok(found_settings) as Result<GraphSettings, String>
+                Ok(()) as Result<(), String>
             })
-            .collect::<Vec<Result<GraphSettings, String>>>();
+            .collect::<Result<Vec<()>, String>>()?;
 
-        let mut res: Vec<GraphSettings> = vec![];
-        for task in tasks {
-            if let Ok(task) = task {
-                res.push(task);
-            }
-        }
+        let min_lookup_range = forward_pass_res
+            .iter()
+            .map(|x| x.1.iter().map(|x| x.min_lookup_inputs))
+            .flatten()
+            .min()
+            .unwrap_or(0);
+
+        let max_lookup_range = forward_pass_res
+            .iter()
+            .map(|x| x.1.iter().map(|x| x.max_lookup_inputs))
+            .flatten()
+            .max()
+            .unwrap_or(0);
+
+        let res = circuit.calibrate_from_min_max(
+            min_lookup_range,
+            max_lookup_range,
+            max_logrows,
+            lookup_safety_margin,
+        );
 
         // drop the gag
         #[cfg(unix)]
@@ -917,31 +900,37 @@ pub(crate) fn calibrate(
         #[cfg(unix)]
         std::mem::drop(_q);
 
-        let max_lookup_range = res
-            .iter()
-            .map(|x| x.run_args.lookup_range.1)
-            .max()
-            .unwrap_or(0);
-        let min_lookup_range = res
-            .iter()
-            .map(|x| x.run_args.lookup_range.0)
-            .min()
-            .unwrap_or(0);
+        if res.is_ok() {
+            let new_settings = circuit.settings().clone();
+
+            let found_run_args = RunArgs {
+                input_scale: new_settings.run_args.input_scale,
+                param_scale: new_settings.run_args.param_scale,
+                lookup_range: new_settings.run_args.lookup_range,
+                logrows: new_settings.run_args.logrows,
+                scale_rebase_multiplier: new_settings.run_args.scale_rebase_multiplier,
+                ..settings.run_args.clone()
+            };
+
+            let found_settings = GraphSettings {
+                run_args: found_run_args,
+                required_lookups: new_settings.required_lookups,
+                model_output_scales: new_settings.model_output_scales,
+                model_input_scales: new_settings.model_input_scales,
+                num_rows: new_settings.num_rows,
+                total_assignments: new_settings.total_assignments,
+                total_const_size: new_settings.total_const_size,
+                ..settings.clone()
+            };
+
+            found_params.push(found_settings.clone());
 
-        if let Some(mut best) = res.into_iter().max_by_key(|p| {
-            (
-                p.run_args.logrows,
-                p.run_args.input_scale,
-                p.run_args.param_scale,
-            )
-        }) {
-            best.run_args.lookup_range = (min_lookup_range, max_lookup_range);
-            // pick the one with the largest logrows
-            found_params.push(best.clone());
             debug!(
                 "found settings: \n {}",
-                best.as_json()?.to_colored_json_auto()?
+                found_settings.as_json()?.to_colored_json_auto()?
             );
+        } else {
+            debug!("calibration failed");
         }
 
         pb.inc(1);

src/graph/mod.rs

@@ -956,10 +956,14 @@ impl GraphCircuit {
         (ASSUMED_BLINDING_FACTORS + RESERVED_BLINDING_ROWS_PAD) as f64
     }
 
-    fn calc_safe_lookup_range(res: &GraphWitness, lookup_safety_margin: i128) -> (i128, i128) {
+    fn calc_safe_lookup_range(
+        min_lookup_inputs: i128,
+        max_lookup_inputs: i128,
+        lookup_safety_margin: i128,
+    ) -> (i128, i128) {
         let mut margin = (
-            lookup_safety_margin * res.min_lookup_inputs,
-            lookup_safety_margin * res.max_lookup_inputs,
+            lookup_safety_margin * min_lookup_inputs,
+            lookup_safety_margin * max_lookup_inputs,
         );
         if lookup_safety_margin == 1 {
             margin.0 -= 1;
@@ -978,7 +982,8 @@ impl GraphCircuit {
 
     fn calc_min_logrows(
         &mut self,
-        res: &GraphWitness,
+        min_lookup_inputs: i128,
+        max_lookup_inputs: i128,
         max_logrows: Option<u32>,
         lookup_safety_margin: i128,
     ) -> Result<(), Box<dyn std::error::Error>> {
@@ -986,19 +991,24 @@ impl GraphCircuit {
         let max_logrows = max_logrows.unwrap_or(MAX_PUBLIC_SRS);
         let max_logrows = std::cmp::min(max_logrows, MAX_PUBLIC_SRS);
         let mut max_logrows = std::cmp::max(max_logrows, MIN_LOGROWS);
+        let mut min_logrows = MIN_LOGROWS;
 
         let reserved_blinding_rows = Self::reserved_blinding_rows();
         // check if has overflowed max lookup input
-        if res.max_lookup_inputs > MAX_LOOKUP_ABS / lookup_safety_margin
-            || res.min_lookup_inputs < -MAX_LOOKUP_ABS / lookup_safety_margin
+        if max_lookup_inputs > MAX_LOOKUP_ABS / lookup_safety_margin
+            || min_lookup_inputs < -MAX_LOOKUP_ABS / lookup_safety_margin
         {
-            let err_string = format!("max lookup input ({}) is too large", res.max_lookup_inputs);
+            let err_string = format!("max lookup input ({}) is too large", max_lookup_inputs);
             error!("{}", err_string);
             return Err(err_string.into());
         }
 
-        let safe_range = Self::calc_safe_lookup_range(res, lookup_safety_margin);
-        let mut min_logrows = MIN_LOGROWS;
+        let safe_range = Self::calc_safe_lookup_range(
+            min_lookup_inputs,
+            max_lookup_inputs,
+            lookup_safety_margin,
+        );
+
         // degrade the max logrows until the extended k is small enough
         while min_logrows < max_logrows
             && !self.extended_k_is_small_enough(
@@ -1020,8 +1030,7 @@ impl GraphCircuit {
             return Err(err_string.into());
         }
 
-        // degrade the max logrows until the extended k is small enough
-        while max_logrows > min_logrows
+        while min_logrows < max_logrows
             && !self.extended_k_is_small_enough(
                 max_logrows,
                 Self::calc_num_cols(safe_range, max_logrows),
@@ -1030,6 +1039,17 @@ impl GraphCircuit {
             max_logrows -= 1;
         }
 
+        if !self
+            .extended_k_is_small_enough(max_logrows, Self::calc_num_cols(safe_range, max_logrows))
+        {
+            let err_string = format!(
+                "extended k is too large to accommodate the quotient polynomial with logrows {}",
+                max_logrows
+            );
+            error!("{}", err_string);
+            return Err(err_string.into());
+        }
+
         let min_bits = ((safe_range.1 - safe_range.0) as f64 + reserved_blinding_rows + 1.)
             .log2()
             .ceil() as usize;
@@ -1111,22 +1131,31 @@ impl GraphCircuit {
         // n = 2^k
         let n = 1u64 << k;
         let mut extended_k = k;
+
         while (1 << extended_k) < (n * quotient_poly_degree) {
             extended_k += 1;
+            if !(extended_k <= bn256::Fr::S) {
+                return false;
+            }
         }
-        extended_k <= bn256::Fr::S
+        true
     }
 
     /// Calibrate the circuit to the supplied data.
-    pub fn calibrate(
+    pub fn calibrate_from_min_max(
         &mut self,
-        input: &[Tensor<Fp>],
+        min_lookup_inputs: i128,
+        max_lookup_inputs: i128,
         max_logrows: Option<u32>,
         lookup_safety_margin: i128,
-    ) -> Result<GraphWitness, Box<dyn std::error::Error>> {
-        let res = self.forward(&mut input.to_vec(), None, None)?;
-        self.calc_min_logrows(&res, max_logrows, lookup_safety_margin)?;
-        Ok(res)
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        self.calc_min_logrows(
+            min_lookup_inputs,
+            max_lookup_inputs,
+            max_logrows,
+            lookup_safety_margin,
+        )?;
+        Ok(())
     }
 
     /// Runs the forward pass of the model / graph of computations and any associated hashing.

tests/integration_tests.rs

@@ -516,7 +516,7 @@ mod native_tests {
                 crate::native_tests::setup_py_env();
                 let test_dir = TempDir::new(test).unwrap();
                 let path = test_dir.path().to_str().unwrap(); crate::native_tests::mv_test_(path, test);
-                accuracy_measurement(path, test.to_string(), "private", "private", "public", 1, "accuracy", 1.2);
+                accuracy_measurement(path, test.to_string(), "private", "private", "public", 1, "accuracy", 2.6);
                 test_dir.close().unwrap();
             }
 
@@ -526,7 +526,7 @@ mod native_tests {
                 crate::native_tests::setup_py_env();
                 let test_dir = TempDir::new(test).unwrap();
                 let path = test_dir.path().to_str().unwrap(); crate::native_tests::mv_test_(path, test);
-                accuracy_measurement(path, test.to_string(), "private", "fixed", "private", 1, "accuracy", 1.2);
+                accuracy_measurement(path, test.to_string(), "private", "fixed", "private", 1, "accuracy", 2.6);
                 test_dir.close().unwrap();
             }
 
@@ -536,7 +536,7 @@ mod native_tests {
                 crate::native_tests::setup_py_env();
                 let test_dir = TempDir::new(test).unwrap();
                 let path = test_dir.path().to_str().unwrap(); crate::native_tests::mv_test_(path, test);
-                accuracy_measurement(path, test.to_string(), "public", "private", "private", 1, "accuracy", 1.2);
+                accuracy_measurement(path, test.to_string(), "public", "private", "private", 1, "accuracy", 2.6);
                 test_dir.close().unwrap();
             }
 
@@ -1901,6 +1901,20 @@ mod native_tests {
                     .map(|h| vec![FileSourceInner::Field(*h)])
                     .collect(),
             ));
+        } else {
+            input.output_data = Some(DataSource::File(
+                witness
+                    .pretty_elements
+                    .unwrap()
+                    .rescaled_outputs
+                    .iter()
+                    .map(|o| {
+                        o.iter()
+                            .map(|f| FileSourceInner::Float(f.parse().unwrap()))
+                            .collect()
+                    })
+                    .collect(),
+            ));
         }
 
         input.save(data_path.clone().into()).unwrap();

tests/py_integration_tests.rs

@@ -118,6 +118,7 @@ mod py_tests {
     }
 
     const TESTS: [&str; 32] = [
+        "variance.ipynb",
         "proof_splitting.ipynb",
         "mnist_gan_proof_splitting.ipynb",
         "mnist_gan.ipynb",
@@ -126,7 +127,6 @@ mod py_tests {
         "hashed_vis.ipynb",
         "simple_demo_all_public.ipynb",
         "data_attest.ipynb",
-        "variance.ipynb",
         "little_transformer.ipynb",
         "simple_demo_aggregated_proofs.ipynb",
         "ezkl_demo.ipynb",