add: leaky_relu non-linearity (zkonduit#72)

Co-authored-by: Alexander Camuto <[email protected]>

.gitignore

@@ -5,3 +5,4 @@ data
 *.params
 *~
 \#*\#
+.DS_Store

Cargo.toml

@@ -26,7 +26,7 @@ foundry_evm = { git = "https://github.com/foundry-rs/foundry", package = "foundr
 halo2_wrong_ecc = { git = "https://github.com/zkonduit/halo2wrong", package = "ecc", branch = "master", optional=true}
 plonk_verifier = { git = "https://github.com/zkonduit/plonk-verifier", branch = "main"}
 colog = { version = "1.1.0", optional = true }
-
+eq-float = "0.1.0"
 
 [dev-dependencies]
 criterion = {version = "0.3",  features = ["html_reports"]}

benches/relu.rs

@@ -1,5 +1,5 @@
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
-use ezkl::circuit::eltwise::{EltwiseConfig, Nonlin1d, Nonlinearity, ReLu};
+use ezkl::circuit::eltwise::{EltwiseConfig, Nonlin1d, Nonlinearity, ReLU};
 use ezkl::tensor::*;
 use halo2_proofs::dev::MockProver;
 use halo2_proofs::{
@@ -37,7 +37,7 @@ impl<F: FieldExt + TensorType, NL: 'static + Nonlinearity<F> + Clone> Circuit<F>
                 .map(|_| VarTensor::new_advice(cs, K, LEN, vec![LEN], true, 512))
                 .collect::<Vec<_>>();
 
-            Self::Config::configure(cs, &advices[0], &advices[1], Some(&[BITS, 128]))
+            Self::Config::configure(cs, &advices[0], &advices[1], BITS, &[128], None)
         }
     }
 
@@ -66,13 +66,13 @@ fn runrelu(c: &mut Criterion) {
         let input: Tensor<Value<F>> =
             Tensor::<i32>::from((0..len).map(|_| rng.gen_range(0..10))).into();
 
-        let assigned: Nonlin1d<F, ReLu<F>> = Nonlin1d {
+        let assigned: Nonlin1d<F, ReLU<F>> = Nonlin1d {
             input: ValTensor::from(input.clone()),
             output: ValTensor::from(input),
             _marker: PhantomData,
         };
 
-        let circuit = NLCircuit::<F, ReLu<F>> {
+        let circuit = NLCircuit::<F, ReLU<F>> {
             assigned,
             _marker: PhantomData,
         };

examples/conv2d_mnist/main.rs

@@ -1,4 +1,4 @@
-use ezkl::circuit::eltwise::{EltwiseConfig, ReLu};
+use ezkl::circuit::eltwise::{EltwiseConfig, ReLU};
 use ezkl::circuit::fused::*;
 use ezkl::fieldutils;
 use ezkl::fieldutils::i32_to_felt;
@@ -54,7 +54,7 @@ struct Config<
 {
     // this will be a conv layer
     l0: FusedConfig<F>,
-    l1: EltwiseConfig<F, ReLu<F>>,
+    l1: EltwiseConfig<F, ReLU<F>>,
     // this will be an affine layer
     l2: FusedConfig<F>,
     public_output: Column<Instance>,
@@ -200,8 +200,8 @@ where
         let input = input.reshape(&[LEN]);
         let output = output.reshape(&[LEN]);
 
-        let l1: EltwiseConfig<F, ReLu<F>> =
-            EltwiseConfig::configure(cs, &input, &output, Some(&[BITS, 32]));
+        let l1: EltwiseConfig<F, ReLU<F>> =
+            EltwiseConfig::configure(cs, &input, &output, BITS, &[32], None);
 
         // tells the config layer to add an affine op to the circuit gate
         let affine_node = FusedNode {
@@ -384,7 +384,7 @@ pub fn runconv() {
         let root = BitMapBackend::new("conv2dmnist-layout.png", (2048, 7680)).into_drawing_area();
         root.fill(&WHITE).unwrap();
         let root = root
-            .titled("Conv -> ReLu -> Affine -> Relu", ("sans-serif", 60))
+            .titled("Conv -> ReLU -> Affine -> ReLU", ("sans-serif", 60))
             .unwrap();
 
         halo2_proofs::dev::CircuitLayout::default()

examples/mlp_4d.rs

@@ -1,4 +1,4 @@
-use ezkl::circuit::eltwise::{DivideBy, EltwiseConfig, ReLu};
+use ezkl::circuit::eltwise::{DivideBy, EltwiseConfig, ReLU};
 use ezkl::circuit::fused::*;
 use ezkl::fieldutils::i32_to_felt;
 use ezkl::tensor::*;
@@ -16,9 +16,9 @@ const K: usize = 15;
 #[derive(Clone)]
 struct MyConfig<F: FieldExt + TensorType> {
     l0: FusedConfig<F>,
-    l1: EltwiseConfig<F, ReLu<F>>,
+    l1: EltwiseConfig<F, ReLU<F>>,
     l2: FusedConfig<F>,
-    l3: EltwiseConfig<F, ReLu<F>>,
+    l3: EltwiseConfig<F, ReLU<F>>,
     l4: EltwiseConfig<F, DivideBy<F>>,
     public_output: Column<Instance>,
 }
@@ -79,12 +79,12 @@ impl<F: FieldExt + TensorType, const LEN: usize, const BITS: usize> Circuit<F>
         );
 
         // sets up a new ReLU table and resuses it for l1 and l3 non linearities
-        let [l1, l3]: [EltwiseConfig<F, ReLu<F>>; 2] =
-            EltwiseConfig::configure_multiple(cs, &input, &output, Some(&[BITS, 1]));
+        let [l1, l3]: [EltwiseConfig<F, ReLU<F>>; 2] =
+            EltwiseConfig::configure_multiple(cs, &input, &output, BITS, &[1], None);
 
         // sets up a new Divide by table
         let l4: EltwiseConfig<F, DivideBy<F>> =
-            EltwiseConfig::configure(cs, &input, &output, Some(&[BITS, 128]));
+            EltwiseConfig::configure(cs, &input, &output, BITS, &[128], None);
 
         let public_output: Column<Instance> = cs.instance_column();
         cs.enable_equality(public_output);

src/circuit/eltwise.rs

@@ -17,10 +17,10 @@ pub trait Nonlinearity<F: FieldExt> {
     ///
     /// * `x` - input to function
     /// * `scales` - additional parameters that may parametrize the function
-    fn nonlinearity(x: i32, scales: &[usize]) -> F;
+    fn nonlinearity(x: i32, scales: &[usize], param: Option<f32>) -> F;
     /// a value which is always in the table
-    fn default_pair(scales: &[usize]) -> (F, F) {
-        (F::zero(), Self::nonlinearity(0, scales))
+    fn default_pair(scales: &[usize], param: Option<f32>) -> (F, F) {
+        (F::zero(), Self::nonlinearity(0, scales, param))
     }
 }
 
@@ -45,8 +45,10 @@ pub struct EltwiseTable<F: FieldExt, NL: Nonlinearity<F>> {
     pub table_output: TableColumn,
     /// Flags if table has been previously assigned to.
     pub is_assigned: bool,
-    /// Number of bits used in lookup table.
+    /// Scaling of the table's inputs.
     pub scaling_params: Vec<usize>,
+    /// Parameters related to the eltwise function being represented
+    pub eltwise_params: Option<f32>,
     /// Number of bits used in lookup table.
     pub bits: usize,
     _marker: PhantomData<(F, NL)>,
@@ -58,12 +60,14 @@ impl<F: FieldExt, NL: Nonlinearity<F>> EltwiseTable<F, NL> {
         cs: &mut ConstraintSystem<F>,
         bits: usize,
         scaling_params: &[usize],
+        eltwise_params: Option<f32>,
     ) -> EltwiseTable<F, NL> {
         EltwiseTable {
             table_input: cs.lookup_table_column(),
             table_output: cs.lookup_table_column(),
             is_assigned: false,
             scaling_params: scaling_params.to_vec(),
+            eltwise_params,
             bits,
             _marker: PhantomData,
         }
@@ -94,7 +98,13 @@ impl<F: FieldExt, NL: Nonlinearity<F>> EltwiseTable<F, NL> {
                         || format!("nl_o_col row {}", row_offset),
                         self.table_output,
                         row_offset,
-                        || Value::known(NL::nonlinearity(int_input, &self.scaling_params)),
+                        || {
+                            Value::known(NL::nonlinearity(
+                                int_input,
+                                &self.scaling_params,
+                                self.eltwise_params,
+                            ))
+                        },
                     ) {
                         Ok(a) => a,
                         Err(e) => {
@@ -133,13 +143,17 @@ impl<F: FieldExt + TensorType, NL: 'static + Nonlinearity<F>> EltwiseConfig<F, N
         cs: &mut ConstraintSystem<F>,
         input: &VarTensor,
         output: &VarTensor,
-        eltwise_params: Option<&[usize]>,
+        bits: usize,
+        scaling_params: &[usize],
+        eltwise_params: Option<f32>,
     ) -> [Self; NUM] {
         let mut table: Option<Rc<RefCell<EltwiseTable<F, NL>>>> = None;
         let configs = (0..NUM)
             .map(|_| {
                 let l = match &table {
-                    None => Self::configure(cs, input, output, eltwise_params),
+                    None => {
+                        Self::configure(cs, input, output, bits, scaling_params, eltwise_params)
+                    }
                     Some(t) => Self::configure_with_table(cs, input, output, t.clone()),
                 };
                 table = Some(l.table.clone());
@@ -168,8 +182,10 @@ impl<F: FieldExt + TensorType, NL: 'static + Nonlinearity<F>> EltwiseConfig<F, N
                 let _ = cs.lookup("lk", |cs| {
                     let qlookup = cs.query_selector(qlookup);
                     let not_qlookup = Expression::Constant(F::one()) - qlookup.clone();
-                    let (default_x, default_y) =
-                        NL::default_pair(table.borrow().scaling_params.as_slice());
+                    let (default_x, default_y) = NL::default_pair(
+                        table.borrow().scaling_params.as_slice(),
+                        table.borrow().eltwise_params,
+                    );
                     let (x, y) = input.cartesian_coord(i);
                     vec![
                         (
@@ -221,20 +237,15 @@ impl<F: FieldExt + TensorType, NL: 'static + Nonlinearity<F>> EltwiseConfig<F, N
         cs: &mut ConstraintSystem<F>,
         input: &VarTensor,
         output: &VarTensor,
-        eltwise_params: Option<&[usize]>,
+        bits: usize,
+        scaling_params: &[usize],
+        eltwise_params: Option<f32>,
     ) -> Self {
-        // will fail if not supplied
-        let params = match eltwise_params {
-            Some(p) => p,
-            None => {
-                panic!("failed to supply eltwise parameters")
-            }
-        };
-        let bits = params[0];
         let table = Rc::new(RefCell::new(EltwiseTable::<F, NL>::configure(
             cs,
             bits,
-            &params[1..],
+            scaling_params,
+            eltwise_params,
         )));
         Self::configure_with_table(cs, input, output, table)
     }
@@ -259,6 +270,7 @@ impl<F: FieldExt + TensorType, NL: 'static + Nonlinearity<F>> EltwiseConfig<F, N
                                 <NL as Nonlinearity<F>>::nonlinearity(
                                     felt_to_i32(f.evaluate()),
                                     &self.table.borrow().scaling_params,
+                                    self.table.borrow().eltwise_params,
                                 )
                             })
                         }));
@@ -283,11 +295,11 @@ impl<F: FieldExt + TensorType, NL: 'static + Nonlinearity<F>> EltwiseConfig<F, N
 #[allow(missing_docs)]
 // Now implement nonlinearity functions like this
 #[derive(Clone, Debug)]
-pub struct ReLu<F> {
+pub struct ReLU<F> {
     _marker: PhantomData<F>,
 }
-impl<F: FieldExt> Nonlinearity<F> for ReLu<F> {
-    fn nonlinearity(x: i32, scale: &[usize]) -> F {
+impl<F: FieldExt> Nonlinearity<F> for ReLU<F> {
+    fn nonlinearity(x: i32, scale: &[usize], _: Option<f32>) -> F {
         if x < 0 {
             F::zero()
         } else {
@@ -298,6 +310,24 @@ impl<F: FieldExt> Nonlinearity<F> for ReLu<F> {
     }
 }
 
+#[allow(missing_docs)]
+#[derive(Clone, Debug)]
+pub struct LeakyReLU<F> {
+    _marker: PhantomData<F>,
+}
+
+impl<F: FieldExt> Nonlinearity<F> for LeakyReLU<F> {
+    fn nonlinearity(x: i32, scale: &[usize], slope: Option<f32>) -> F {
+        if x < 0 {
+            let d_inv_x = slope.unwrap() * (x as f32) / (scale[0] as f32);
+            let rounded = d_inv_x.round();
+            fieldutils::i32_to_felt(rounded as i32)
+        } else {
+            fieldutils::i32_to_felt(x)
+        }
+    }
+}
+
 #[allow(missing_docs)]
 #[derive(Clone, Debug)]
 pub struct Sigmoid<F> {
@@ -306,7 +336,7 @@ pub struct Sigmoid<F> {
 // L is our implicit or explicit denominator (fixed point d)
 // Usually want K=L
 impl<F: FieldExt> Nonlinearity<F> for Sigmoid<F> {
-    fn nonlinearity(x: i32, scale: &[usize]) -> F {
+    fn nonlinearity(x: i32, scale: &[usize], _: Option<f32>) -> F {
         let kix = (x as f32) / (scale[0] as f32);
         let fout = (scale[1] as f32) / (1.0 + (-kix).exp());
         let rounded = fout.round();
@@ -320,7 +350,7 @@ pub struct DivideBy<F> {
     _marker: PhantomData<F>,
 }
 impl<F: FieldExt> Nonlinearity<F> for DivideBy<F> {
-    fn nonlinearity(x: i32, scale: &[usize]) -> F {
+    fn nonlinearity(x: i32, scale: &[usize], _: Option<f32>) -> F {
         let d_inv_x = (x as f32) / (scale[0] as f32);
         let rounded = d_inv_x.round();
         fieldutils::i32_to_felt(rounded as i32)
@@ -359,7 +389,7 @@ mod tests {
                 .map(|_| VarTensor::new_advice(cs, 4, 1, vec![1], true, 512))
                 .collect::<Vec<_>>();
 
-            Self::Config::configure(cs, &advices[0], &advices[1], Some(&[2, 1]))
+            Self::Config::configure(cs, &advices[0], &advices[1], 2, &[1], None)
         }
 
         fn synthesize(
@@ -376,33 +406,46 @@ mod tests {
     #[test]
     fn test_eltrelunl() {
         for i in -127..127 {
-            let r = <ReLu<F> as Nonlinearity<F>>::nonlinearity(i, &[1]);
+            let r = <ReLU<F> as Nonlinearity<F>>::nonlinearity(i, &[1], None);
+            if i <= 0 {
+                assert_eq!(r, F::from(0_u64))
+            } else {
+                assert_eq!(r, F::from(i as u64))
+            }
+        }
+    }
+
+    #[test]
+    fn test_eltleakyrelunl() {
+        for i in -127..127 {
+            let r = <LeakyReLU<F> as Nonlinearity<F>>::nonlinearity(i, &[1], Some(0.05));
             if i <= 0 {
-                assert!(r == F::from(0_u64))
+                println!("{:?}", (0.05 * i as f32));
+                assert_eq!(r, -F::from(-(0.05 * i as f32).round() as u64))
             } else {
-                assert!(r == F::from(i as u64))
+                assert_eq!(r, F::from(i as u64))
             }
         }
     }
 
     #[test]
     fn test_eltsigmoid() {
         for i in -127..127 {
-            let r = <Sigmoid<F> as Nonlinearity<F>>::nonlinearity(i, &[1, 1]);
+            let r = <Sigmoid<F> as Nonlinearity<F>>::nonlinearity(i, &[1, 1], None);
             let exp_sig = (1.0 / (1.0 + (-i as f32).exp())).round();
-            assert!(r == F::from(exp_sig as u64))
+            assert_eq!(r, F::from(exp_sig as u64))
         }
     }
 
     #[test]
     fn test_eltdivide() {
         for i in -127..127 {
-            let r = <DivideBy<F> as Nonlinearity<F>>::nonlinearity(i, &[1]);
+            let r = <DivideBy<F> as Nonlinearity<F>>::nonlinearity(i, &[1], None);
             println!("{:?}, {:?}, {:?}", i, r, F::from(-i as u64));
             if i <= 0 {
-                assert!(r == -F::from(-i as u64))
+                assert_eq!(r, -F::from(-i as u64))
             } else {
-                assert!(r == F::from(i as u64))
+                assert_eq!(r, F::from(i as u64))
             }
         }
     }
@@ -411,13 +454,13 @@ mod tests {
     fn relucircuit() {
         let input: Tensor<Value<F>> =
             Tensor::new(Some(&[Value::<F>::known(F::from(1_u64))]), &[1]).unwrap();
-        let assigned: Nonlin1d<F, ReLu<F>> = Nonlin1d {
+        let assigned: Nonlin1d<F, ReLU<F>> = Nonlin1d {
             input: ValTensor::from(input.clone()),
             output: ValTensor::from(input),
             _marker: PhantomData,
         };
 
-        let circuit = NLCircuit::<F, ReLu<F>> {
+        let circuit = NLCircuit::<F, ReLU<F>> {
             assigned,
             _marker: PhantomData,
         };