Dimitris/requantize bmm ref

README.md

@@ -18,13 +18,13 @@ where `<example_name>` is one of the following:
 
 In order to run any tests involving python code, such as compatibility tests with TF Lite, the feature `python` must be activated (which automatically enables `test-types`).
 
-## From `ndarray` to `QArray`
+## From `ndarray` to `Tensor`
 
-In order to save a `numpy` `ndarray` (python side) as a serialised JSON which can be directly read into a `QArray` of ours (Rust side),
+In order to save a `numpy` `ndarray` (python side) as a serialised JSON which can be directly read into a `Tensor` of ours (Rust side),
 - Convert the `ndarray` into an `OrderedDict` using our custom python function `tensor_to_dict` (available in several of the python notebooks)
 - Pass the resulting `OrderedDict` together with the destination path to `json.dump`.
 
-The saved JSON file can be deserialised over in Rust with `QArray::read(path: &str) -> QArray`. If instead of a single `OrderedDict`, a python list of `OrderedDict`s is passed to `json.dump`, the resulting file can be deserialised with `QArray::read_list(path: &str) -> Vec<QArray> `.
+The saved JSON file can be deserialised over in Rust with `Tensor::read(path: &str) -> Tensor`. If instead of a single `OrderedDict`, a python list of `OrderedDict`s is passed to `json.dump`, the resulting file can be deserialised with `Tensor::read_list(path: &str) -> Vec<Tensor> `.
 
 Cf. `exploring_tf_lite/training_two_layer_perceptron.ipynb` for example usage.
 

common/examples/common/lib.rs

@@ -1,21 +1,21 @@
-use hcs_common::{quantise_f32_u8_nne, Model, Poly, QArray};
+use hcs_common::{quantise_f32_u8_nne, Model, Poly, Tensor};
 
 use ark_crypto_primitives::sponge::{Absorb, CryptographicSponge};
 use ark_ff::PrimeField;
 use ark_poly_commit::PolynomialCommitment;
 
 // Auxiliary function
 fn unpadded_inference<F, S, PCS>(
-    raw_input: QArray<f32>,
-    model: &Model<i8, i32>,
+    raw_input: Tensor<f32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
-) -> QArray<u8>
+) -> Tensor<u8>
 where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let quantised_input: QArray<u8> = QArray::new(
+    let quantised_input: Tensor<u8> = Tensor::new(
         quantise_f32_u8_nne(raw_input.values(), qinfo.0, qinfo.1),
         raw_input.shape().clone(),
     );
@@ -31,40 +31,40 @@ where
 // If padded inference is left on the prover side, move this to the prover
 /* // Auxiliary function
 fn padded_inference<F, S, PCS>(
-    raw_input: QArray<f32>,
-    model: &Model<i8, i32>,
+    raw_input: Tensor<f32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
-) -> QArray<u8>
+) -> Tensor<u8>
 where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let quantised_input: QArray<u8> = QArray::new(
+    let quantised_input: Tensor<u8> = Tensor::new(
         quantise_f32_u8_nne(raw_input.values(), qinfo.0, qinfo.1),
         raw_input.shape().clone(),
     );
 
     let input_i8 = (quantised_input.cast::<i32>() - 128).cast::<i8>();
 
     let output_i8 =
-        <Model<i8, i32> as ProveModel<F, S, PCS, i8, i32>>::padded_evaluate(model, input_i8);
+        <Model<i8> as ProveModel<F, S, PCS, i8>>::padded_evaluate(model, input_i8);
 
     (output_i8.cast::<i32>() + 128).cast()
 } */
 
 pub fn run_unpadded<F, S, PCS>(
     input_path: &str,
     expected_output_path: &str,
-    model: &Model<i8, i32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
 ) where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let raw_input: QArray<f32> = QArray::read(input_path);
-    let expected_output: QArray<u8> = QArray::read(expected_output_path);
+    let raw_input: Tensor<f32> = Tensor::read(input_path);
+    let expected_output: Tensor<u8> = Tensor::read(expected_output_path);
 
     let output_u8 = unpadded_inference::<F, S, PCS>(raw_input, model, qinfo);
 
@@ -78,15 +78,15 @@ pub fn run_unpadded<F, S, PCS>(
 /* pub fn run_padded<F, S, PCS>(
     input_path: &str,
     expected_output_path: &str,
-    model: &Model<i8, i32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
 ) where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let raw_input: QArray<f32> = QArray::read(input_path);
-    let expected_output: QArray<u8> = QArray::read(expected_output_path);
+    let raw_input: Tensor<f32> = Tensor::read(input_path);
+    let expected_output: Tensor<u8> = Tensor::read(expected_output_path);
 
     let output_u8 = padded_inference::<F, S, PCS>(raw_input, model, qinfo);
 
@@ -98,15 +98,15 @@ pub fn run_unpadded<F, S, PCS>(
 pub fn multi_run_unpadded<F, S, PCS>(
     inputs_path: &str,
     expected_outputs_path: &str,
-    model: &Model<i8, i32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
 ) where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let raw_inputs: Vec<QArray<f32>> = QArray::read_list(inputs_path);
-    let expected_outputs: Vec<QArray<u8>> = QArray::read_list(expected_outputs_path);
+    let raw_inputs: Vec<Tensor<f32>> = Tensor::read_list(inputs_path);
+    let expected_outputs: Vec<Tensor<u8>> = Tensor::read_list(expected_outputs_path);
 
     for (raw_input, expected_output) in raw_inputs.into_iter().zip(expected_outputs.into_iter()) {
         assert_eq!(
@@ -131,8 +131,8 @@ pub fn multi_run_padded<F, S, PCS>(
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let raw_inputs: Vec<QArray<f32>> = QArray::read_list(inputs_path);
-    let expected_outputs: Vec<QArray<u8>> = QArray::read_list(expected_outputs_path);
+    let raw_inputs: Vec<Tensor<f32>> = Tensor::read_list(inputs_path);
+    let expected_outputs: Vec<Tensor<u8>> = Tensor::read_list(expected_outputs_path);
 
     for (raw_input, expected_output) in raw_inputs.into_iter().zip(expected_outputs.into_iter()) {
         assert_eq!(

common/examples/simple_perceptron_mnist/main.rs

@@ -1,6 +1,6 @@
 use hcs_common::{
     simple_perceptron_mnist::{build_simple_perceptron_mnist, parameters::*},
-    Ligero,
+    BMMRequantizationStrategy, Ligero,
 };
 
 use ark_bn254::Fr;
@@ -20,7 +20,9 @@ macro_rules! PATH {
 }
 
 fn main() {
-    let simple_perceptron = build_simple_perceptron_mnist::<Fr, PoseidonSponge<Fr>, Ligero<Fr>>();
+    let simple_perceptron = build_simple_perceptron_mnist::<Fr, PoseidonSponge<Fr>, Ligero<Fr>>(
+        BMMRequantizationStrategy::Floating,
+    );
 
     // Right now this can't be QInfo because the latter is always a pair
     // (f32, i8), which indeed matches in-model quantisation, but not

common/examples/two_layer_perceptron_mnist/main.rs

@@ -1,6 +1,6 @@
 use hcs_common::{
     two_layer_perceptron_mnist::{build_two_layer_perceptron_mnist, parameters::*},
-    Ligero,
+    BMMRequantizationStrategy, Ligero,
 };
 
 use ark_bn254::Fr;
@@ -20,8 +20,9 @@ macro_rules! PATH {
 }
 
 fn main() {
-    let two_layer_perceptron =
-        build_two_layer_perceptron_mnist::<Fr, PoseidonSponge<Fr>, Ligero<Fr>>();
+    let two_layer_perceptron = build_two_layer_perceptron_mnist::<Fr, PoseidonSponge<Fr>, Ligero<Fr>>(
+        BMMRequantizationStrategy::Floating,
+    );
 
     // Right now this can't be QInfo because the latter is always a pair
     // (f32, i8), which indeed matches in-model quantisation, but not

common/src/compatibility/example_models/simple_perceptron_mnist/mod.rs

@@ -1,4 +1,7 @@
-use crate::{BMMNode, Model, Node, Poly, QArray, RequantiseBMMNode, ReshapeNode};
+use crate::{
+    quantization::BMMRequantizationStrategy, utils::req_bmm_from_strategy, BMMNode, Model, Node,
+    Poly, ReshapeNode, Tensor,
+};
 
 use ark_crypto_primitives::sponge::{Absorb, CryptographicSponge};
 use ark_ff::PrimeField;
@@ -22,7 +25,9 @@ macro_rules! PATH {
 }
 
 // TODO this is incorrect now that we have switched to logs
-pub fn build_simple_perceptron_mnist<F, S, PCS>() -> Model<i8, i32>
+pub fn build_simple_perceptron_mnist<F, S, PCS>(
+    req_strategy: BMMRequantizationStrategy,
+) -> Model<i8>
 where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
@@ -32,20 +37,15 @@ where
 
     let reshape: ReshapeNode = ReshapeNode::new(INPUT_DIMS.to_vec(), vec![flat_dim]);
 
-    let w_array: QArray<i8> = QArray::read(&format!(PATH!(), "weights.json"));
-    let b_array: QArray<i32> = QArray::read(&format!(PATH!(), "bias.json"));
+    let w_array: Tensor<i8> = Tensor::read(&format!(PATH!(), "weights.json"));
+    let b_array: Tensor<i32> = Tensor::read(&format!(PATH!(), "bias.json"));
 
-    let bmm: BMMNode<i8, i32> = BMMNode::new(w_array, b_array, Z_I);
+    let bmm: BMMNode<i8> = BMMNode::new(w_array, b_array, Z_I);
 
-    let req_bmm: RequantiseBMMNode<i8> =
-        RequantiseBMMNode::new(OUTPUT_DIM, S_I, Z_I, S_W, Z_W, S_O, Z_O);
+    let req_bmm = req_bmm_from_strategy(req_strategy, OUTPUT_DIM, S_I, Z_I, S_W, Z_W, S_O, Z_O);
 
     Model::new(
         INPUT_DIMS.to_vec(),
-        vec![
-            Node::Reshape(reshape),
-            Node::BMM(bmm),
-            Node::RequantiseBMM(req_bmm),
-        ],
+        vec![Node::Reshape(reshape), Node::BMM(bmm), req_bmm],
     )
 }

common/src/compatibility/example_models/two_layer_perceptron_mnist/mod.rs

@@ -5,7 +5,10 @@ use ark_poly_commit::PolynomialCommitment;
 pub mod parameters;
 use parameters::*;
 
-use crate::{BMMNode, Model, Node, Poly, QArray, ReLUNode, RequantiseBMMNode, ReshapeNode};
+use crate::{
+    quantization::BMMRequantizationStrategy, utils::req_bmm_from_strategy, BMMNode, Model, Node,
+    Poly, ReLUNode, ReshapeNode, Tensor,
+};
 
 pub const INPUT_DIMS: &[usize] = &[28, 28];
 pub const INTER_DIM: usize = 28;
@@ -22,7 +25,9 @@ macro_rules! PATH {
     };
 }
 
-pub fn build_two_layer_perceptron_mnist<F, S, PCS>() -> Model<i8, i32>
+pub fn build_two_layer_perceptron_mnist<F, S, PCS>(
+    req_strategy: BMMRequantizationStrategy,
+) -> Model<i8>
 where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
@@ -32,32 +37,48 @@ where
 
     let reshape: ReshapeNode = ReshapeNode::new(INPUT_DIMS.to_vec(), vec![flat_dim]);
 
-    let w1_array: QArray<i8> = QArray::read(&format!(PATH!(), "weights_1.json"));
-    let b1_array: QArray<i32> = QArray::read(&format!(PATH!(), "bias_1.json"));
-    let w2_array: QArray<i8> = QArray::read(&format!(PATH!(), "weights_2.json"));
-    let b2_array: QArray<i32> = QArray::read(&format!(PATH!(), "bias_2.json"));
+    let w1_array: Tensor<i8> = Tensor::read(&format!(PATH!(), "weights_1.json"));
+    let b1_array: Tensor<i32> = Tensor::read(&format!(PATH!(), "bias_1.json"));
+    let w2_array: Tensor<i8> = Tensor::read(&format!(PATH!(), "weights_2.json"));
+    let b2_array: Tensor<i32> = Tensor::read(&format!(PATH!(), "bias_2.json"));
 
-    let bmm_1: BMMNode<i8, i32> = BMMNode::new(w1_array, b1_array, Z_1_I);
+    let bmm_1: BMMNode<i8> = BMMNode::new(w1_array, b1_array, Z_1_I);
 
-    let req_bmm_1: RequantiseBMMNode<i8> =
-        RequantiseBMMNode::new(INTER_DIM, S_1_I, Z_1_I, S_1_W, Z_1_W, S_1_O, Z_1_O);
+    let req_bmm_1 = req_bmm_from_strategy(
+        req_strategy,
+        INTER_DIM,
+        S_1_I,
+        Z_1_I,
+        S_1_W,
+        Z_1_W,
+        S_1_O,
+        Z_1_O,
+    );
 
     let relu: ReLUNode<i8> = ReLUNode::new(28, Z_1_O);
 
-    let bmm_2: BMMNode<i8, i32> = BMMNode::new(w2_array, b2_array, Z_2_I);
+    let bmm_2: BMMNode<i8> = BMMNode::new(w2_array, b2_array, Z_2_I);
 
-    let req_bmm_2: RequantiseBMMNode<i8> =
-        RequantiseBMMNode::new(OUTPUT_DIM, S_2_I, Z_2_I, S_2_W, Z_2_W, S_2_O, Z_2_O);
+    let req_bmm_2 = req_bmm_from_strategy(
+        req_strategy,
+        OUTPUT_DIM,
+        S_2_I,
+        Z_2_I,
+        S_2_W,
+        Z_2_W,
+        S_2_O,
+        Z_2_O,
+    );
 
     Model::new(
         INPUT_DIMS.to_vec(),
         vec![
             Node::Reshape(reshape),
             Node::BMM(bmm_1),
-            Node::RequantiseBMM(req_bmm_1),
+            req_bmm_1,
             Node::ReLU(relu),
             Node::BMM(bmm_2),
-            Node::RequantiseBMM(req_bmm_2),
+            req_bmm_2,
         ],
     )
 }

common/src/compatibility/python/mod.rs

@@ -5,7 +5,7 @@ use std::{fs::create_dir_all, path::Path};
 
 use pyo3::{prelude::*, PyAny};
 
-use crate::QArray;
+use crate::Tensor;
 
 const PERCEPTRON_PATH: &str = include_str!(concat!(
     env!("CARGO_MANIFEST_DIR"),
@@ -21,7 +21,7 @@ pub fn get_model(py: Python, model_name: &str, args: Option<Vec<(&str, &str)>>)
     func.call1(py, (model_name, args)).unwrap()
 }
 
-pub fn save_model_parameters_as_qarray(py: Python, model: &Py<PyAny>, path: &str) {
+pub fn save_model_parameters_as_tensor(py: Python, model: &Py<PyAny>, path: &str) {
     let path = Path::new(path);
 
     if !path.exists() {
@@ -30,13 +30,13 @@ pub fn save_model_parameters_as_qarray(py: Python, model: &Py<PyAny>, path: &str
     }
 
     model
-        .call_method1(py, "save_params_as_qarray", (path,))
+        .call_method1(py, "save_params_as_verifiaml_tensor", (path,))
         .unwrap();
 }
 
-pub fn get_model_input<'py, T>(python: Python<'py>, model: &Py<PyAny>, index: usize) -> QArray<f32>
+pub fn get_model_input<'py, T>(python: Python<'py>, model: &Py<PyAny>, index: usize) -> Tensor<f32>
 where
-    T: Into<QArray<f32>> + FromPyObject<'py> + Clone,
+    T: Into<Tensor<f32>> + FromPyObject<'py> + Clone,
 {
     let result = model.call_method1(python, "get_input", (index,));
 
@@ -46,10 +46,10 @@ where
     model_input.into()
 }
 
-pub fn get_model_output(py: Python, model: &Py<PyAny>, index: usize) -> QArray<u8> {
+pub fn get_model_output(py: Python, model: &Py<PyAny>, index: usize) -> Tensor<u8> {
     let result = model.call_method1(py, "get_output", (index,));
     // Downcast the result to the expected type
     let model_output = result.unwrap().extract::<Vec<u8>>(py).unwrap();
 
-    QArray::from(model_output)
+    Tensor::from(model_output)
 }