ci: update version string in docs

Cargo.toml

@@ -169,7 +169,7 @@ harness = false
 
 
 [[bench]]
-name = "relu"
+name = "sigmoid"
 harness = false
 
 [[bench]]
@@ -177,12 +177,12 @@ name = "relu_lookupless"
 harness = false
 
 [[bench]]
-name = "accum_matmul_relu"
+name = "accum_matmul_sigmoid"
 harness = false
 
 
 [[bench]]
-name = "accum_matmul_relu_overflow"
+name = "accum_matmul_sigmoid_overflow"
 harness = false
 
 [[bin]]

benches/accum_matmul_sigmoid.rs

@@ -64,7 +64,7 @@ impl Circuit<Fr> for MyCircuit {
                 &a,
                 BITS,
                 K,
-                &LookupOp::LeakyReLU { slope: 0.0.into() },
+                &LookupOp::Sigmoid { scale: 1.0.into() },
             )
             .unwrap();
 
@@ -93,7 +93,7 @@ impl Circuit<Fr> for MyCircuit {
                     .layout(
                         &mut region,
                         &[output.unwrap()],
-                        Box::new(LookupOp::LeakyReLU { slope: 0.0.into() }),
+                        Box::new(LookupOp::Sigmoid { scale: 1.0.into() }),
                     )
                     .unwrap();
                 Ok(())

benches/accum_matmul_sigmoid_overflow.rs

@@ -65,7 +65,7 @@ impl Circuit<Fr> for MyCircuit {
                 &a,
                 BITS,
                 k,
-                &LookupOp::LeakyReLU { slope: 0.0.into() },
+                &LookupOp::Sigmoid { scale: 1.0.into() },
             )
             .unwrap();
 
@@ -94,7 +94,7 @@ impl Circuit<Fr> for MyCircuit {
                     .layout(
                         &mut region,
                         &[output.unwrap()],
-                        Box::new(LookupOp::LeakyReLU { slope: 0.0.into() }),
+                        Box::new(LookupOp::Sigmoid { scale: 1.0.into() }),
                     )
                     .unwrap();
                 Ok(())

benches/relu_lookupless.rs

@@ -68,7 +68,14 @@ impl Circuit<Fr> for NLCircuit {
             |region| {
                 let mut region = RegionCtx::new(region, 0, 1, 1024, 2);
                 config
-                    .layout(&mut region, &[self.input.clone()], Box::new(PolyOp::ReLU))
+                    .layout(
+                        &mut region,
+                        &[self.input.clone()],
+                        Box::new(PolyOp::LeakyReLU {
+                            slope: 0.0.into(),
+                            scale: 1,
+                        }),
+                    )
                     .unwrap();
                 Ok(())
             },

benches/sigmoid.rs

@@ -42,7 +42,7 @@ impl Circuit<Fr> for NLCircuit {
                 .map(|_| VarTensor::new_advice(cs, K, 1, LEN))
                 .collect::<Vec<_>>();
 
-            let nl = LookupOp::LeakyReLU { slope: 0.0.into() };
+            let nl = LookupOp::Sigmoid { scale: 1.0.into() };
 
             let mut config = Config::default();
 
@@ -68,7 +68,7 @@ impl Circuit<Fr> for NLCircuit {
                     .layout(
                         &mut region,
                         &[self.input.clone()],
-                        Box::new(LookupOp::LeakyReLU { slope: 0.0.into() }),
+                        Box::new(LookupOp::Sigmoid { scale: 1.0.into() }),
                     )
                     .unwrap();
                 Ok(())

docs/python/requirements-docs.txt

@@ -1,4 +1,4 @@
-ezkl==0.0.0
+ezkl==15.1.2
 sphinx
 sphinx-rtd-theme
 sphinxcontrib-napoleon

docs/python/src/conf.py

@@ -1,7 +1,7 @@
 import ezkl
 
 project = 'ezkl'
-release = '0.0.0'
+release = '15.1.2'
 version = release
 
 

examples/conv2d_mnist/main.rs

@@ -146,6 +146,8 @@ where
         let params = VarTensor::new_advice(cs, K, NUM_INNER_COLS, LEN);
         let output = VarTensor::new_advice(cs, K, NUM_INNER_COLS, LEN);
 
+        let _constant = VarTensor::constant_cols(cs, K, LEN, false);
+
         println!("INPUT COL {:#?}", input);
 
         let mut layer_config = PolyConfig::configure(
@@ -156,15 +158,11 @@ where
         );
 
         layer_config
-            .configure_lookup(
-                cs,
-                &input,
-                &output,
-                &params,
-                (LOOKUP_MIN, LOOKUP_MAX),
-                K,
-                &LookupOp::LeakyReLU { slope: 0.0.into() },
-            )
+            .configure_range_check(cs, &input, &params, (-1, 1), K)
+            .unwrap();
+
+        layer_config
+            .configure_range_check(cs, &input, &params, (0, 1023), K)
             .unwrap();
 
         layer_config
@@ -195,6 +193,11 @@ where
     ) -> Result<(), Error> {
         config.layer_config.layout_tables(&mut layouter).unwrap();
 
+        config
+            .layer_config
+            .layout_range_checks(&mut layouter)
+            .unwrap();
+
         let x = layouter
             .assign_region(
                 || "mlp_4d",
@@ -224,7 +227,10 @@ where
                         .layout(
                             &mut region,
                             &[x.unwrap()],
-                            Box::new(LookupOp::LeakyReLU { slope: 0.0.into() }),
+                            Box::new(PolyOp::LeakyReLU {
+                                slope: 0.0.into(),
+                                scale: 1,
+                            }),
                         )
                         .unwrap();
 

examples/mlp_4d_einsum.rs

@@ -53,6 +53,10 @@ impl<const LEN: usize, const LOOKUP_MIN: IntegerRep, const LOOKUP_MAX: IntegerRe
         let output = VarTensor::new_advice(cs, K, 1, LEN);
         // tells the config layer to add an affine op to the circuit gate
 
+        let _constant = VarTensor::constant_cols(cs, K, LEN, false);
+
+        println!("INPUT COL {:#?}", input);
+
         let mut layer_config = PolyConfig::<F>::configure(
             cs,
             &[input.clone(), params.clone()],
@@ -60,17 +64,12 @@ impl<const LEN: usize, const LOOKUP_MIN: IntegerRep, const LOOKUP_MAX: IntegerRe
             CheckMode::SAFE,
         );
 
-        // sets up a new ReLU table and resuses it for l1 and l3 non linearities
         layer_config
-            .configure_lookup(
-                cs,
-                &input,
-                &output,
-                &params,
-                (LOOKUP_MIN, LOOKUP_MAX),
-                K,
-                &LookupOp::LeakyReLU { slope: 0.0.into() },
-            )
+            .configure_range_check(cs, &input, &params, (-1, 1), K)
+            .unwrap();
+
+        layer_config
+            .configure_range_check(cs, &input, &params, (0, 1023), K)
             .unwrap();
 
         // sets up a new ReLU table and resuses it for l1 and l3 non linearities
@@ -104,6 +103,11 @@ impl<const LEN: usize, const LOOKUP_MIN: IntegerRep, const LOOKUP_MAX: IntegerRe
     ) -> Result<(), Error> {
         config.layer_config.layout_tables(&mut layouter).unwrap();
 
+        config
+            .layer_config
+            .layout_range_checks(&mut layouter)
+            .unwrap();
+
         let x = layouter
             .assign_region(
                 || "mlp_4d",
@@ -144,7 +148,10 @@ impl<const LEN: usize, const LOOKUP_MIN: IntegerRep, const LOOKUP_MAX: IntegerRe
                         .layout(
                             &mut region,
                             &[x],
-                            Box::new(LookupOp::LeakyReLU { slope: 0.0.into() }),
+                            Box::new(PolyOp::LeakyReLU {
+                                scale: 1,
+                                slope: 0.0.into(),
+                            }),
                         )
                         .unwrap()
                         .unwrap();
@@ -184,7 +191,10 @@ impl<const LEN: usize, const LOOKUP_MIN: IntegerRep, const LOOKUP_MAX: IntegerRe
                         .layout(
                             &mut region,
                             &[x],
-                            Box::new(LookupOp::LeakyReLU { slope: 0.0.into() }),
+                            Box::new(PolyOp::LeakyReLU {
+                                scale: 1,
+                                slope: 0.0.into(),
+                            }),
                         )
                         .unwrap();
                     println!("6");

examples/onnx/rsqrt/gen.py

@@ -0,0 +1,42 @@
+from torch import nn
+import torch
+import json
+import numpy as np
+
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+
+    def forward(self, x):         
+        # reciprocal sqrt 
+        m = 1 / torch.sqrt(x)
+        return m 
+
+
+circuit = MyModel()
+
+x = torch.empty(1, 8).uniform_(0, 1)
+
+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/rsqrt/input.json

@@ -0,0 +1 @@
+{"input_data": [[0.8590779900550842, 0.4029041528701782, 0.6507361531257629, 0.9782488942146301, 0.37392884492874146, 0.6867020726203918, 0.11407750844955444, 0.362740159034729]]}

examples/onnx/rsqrt/network.onnx

@@ -0,0 +1,17 @@
+pytorch2.2.2:Ź

+$
+input/Sqrt_output_0/Sqrt"Sqrt
+1
+/Sqrt_output_0output/Reciprocal"
+Reciprocal
+main_graphZ!
+input
+

+
+batch_size
+b"
+output
+

+
+batch_size
+B

src/circuit/ops/errors.rs

@@ -94,4 +94,7 @@ pub enum CircuitError {
     #[error("[io] {0}")]
     /// IO error
     IoError(#[from] std::io::Error),
+    /// Invalid scale
+    #[error("negative scale for an op that requires positive inputs {0}")]
+    NegativeScale(String),
 }

src/circuit/ops/hybrid.rs

@@ -16,7 +16,6 @@ pub enum HybridOp {
     Recip {
         input_scale: utils::F32,
         output_scale: utils::F32,
-        use_range_check_for_int: bool,
     },
     Div {
         denom: utils::F32,
@@ -45,6 +44,8 @@ pub enum HybridOp {
     ReduceArgMin {
         dim: usize,
     },
+    Max,
+    Min,
     Softmax {
         input_scale: utils::F32,
         output_scale: utils::F32,
@@ -79,6 +80,8 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Hybrid
             | HybridOp::Less { .. }
             | HybridOp::Equals { .. }
             | HybridOp::GreaterEqual { .. }
+            | HybridOp::Max
+            | HybridOp::Min
             | HybridOp::LessEqual { .. } => {
                 vec![0, 1]
             }
@@ -93,13 +96,14 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Hybrid
 
     fn as_string(&self) -> String {
         match self {
+            HybridOp::Max => format!("MAX"),
+            HybridOp::Min => format!("MIN"),
             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
+                "RECIP (input_scale={}, output_scale={})",
+                input_scale, output_scale
             ),
             HybridOp::Div {
                 denom,
@@ -162,6 +166,8 @@ 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::Max => layouts::max_comp(config, region, values[..].try_into()?)?,
+            HybridOp::Min => layouts::min_comp(config, region, values[..].try_into()?)?,
             HybridOp::SumPool {
                 padding,
                 stride,
@@ -179,31 +185,13 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Hybrid
             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()?,
-                        integer_rep_to_felt(input_scale.0 as i128),
-                        integer_rep_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,
-                        },
-                    )?
-                }
-            }
+            } => layouts::recip(
+                config,
+                region,
+                values[..].try_into()?,
+                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,

src/circuit/ops/layouts.rs

@@ -4155,6 +4155,48 @@ pub(crate) fn argmin<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
     Ok(assigned_argmin)
 }
 
+/// max layout
+pub(crate) fn max_comp<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
+    config: &BaseConfig<F>,
+    region: &mut RegionCtx<F>,
+    values: &[ValTensor<F>; 2],
+) -> Result<ValTensor<F>, CircuitError> {
+    let is_greater = greater(config, region, values)?;
+    let is_less = not(config, region, &[is_greater.clone()])?;
+
+    let max_val_p1 = pairwise(
+        config,
+        region,
+        &[values[0].clone(), is_greater],
+        BaseOp::Mult,
+    )?;
+
+    let max_val_p2 = pairwise(config, region, &[values[1].clone(), is_less], BaseOp::Mult)?;
+
+    pairwise(config, region, &[max_val_p1, max_val_p2], BaseOp::Add)
+}
+
+/// min comp layout
+pub(crate) fn min_comp<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
+    config: &BaseConfig<F>,
+    region: &mut RegionCtx<F>,
+    values: &[ValTensor<F>; 2],
+) -> Result<ValTensor<F>, CircuitError> {
+    let is_greater = greater(config, region, values)?;
+    let is_less = not(config, region, &[is_greater.clone()])?;
+
+    let min_val_p1 = pairwise(config, region, &[values[0].clone(), is_less], BaseOp::Mult)?;
+
+    let min_val_p2 = pairwise(
+        config,
+        region,
+        &[values[1].clone(), is_greater],
+        BaseOp::Mult,
+    )?;
+
+    pairwise(config, region, &[min_val_p1, min_val_p2], BaseOp::Add)
+}
+
 /// max layout
 pub(crate) fn max<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
     config: &BaseConfig<F>,
@@ -4263,7 +4305,6 @@ pub(crate) fn sign<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
 ) -> Result<ValTensor<F>, CircuitError> {
     let mut decomp = decompose(config, region, values, &region.base(), &region.legs())?;
     // get every n elements now, which correspond to the sign bit
-
     decomp.get_every_n(region.legs() + 1)?;
     decomp.reshape(values[0].dims())?;
 
@@ -4280,10 +4321,12 @@ pub(crate) fn abs<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
     pairwise(config, region, &[values[0].clone(), sign], BaseOp::Mult)
 }
 
-pub(crate) fn relu<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
+pub(crate) fn leaky_relu<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
     config: &BaseConfig<F>,
     region: &mut RegionCtx<F>,
     values: &[ValTensor<F>; 1],
+    alpha: &utils::F32,
+    input_scale: &i32,
 ) -> Result<ValTensor<F>, CircuitError> {
     let sign = sign(config, region, values)?;
 
@@ -4292,12 +4335,45 @@ pub(crate) fn relu<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
 
     let relu_mask = equals(config, region, &[sign, unit])?;
 
-    pairwise(
+    let positive = pairwise(
         config,
         region,
-        &[values[0].clone(), relu_mask],
+        &[values[0].clone(), relu_mask.clone()],
         BaseOp::Mult,
-    )
+    )?;
+
+    if alpha.0 == 0. {
+        return Ok(positive);
+    }
+
+    if input_scale < &0 {
+        return Err(CircuitError::NegativeScale("leaky_relu".to_string()));
+    }
+
+    let scale_constant = create_constant_tensor(F::from(2_i32.pow(*input_scale as u32) as u64), 1);
+
+    let rescaled_positive = pairwise(config, region, &[positive, scale_constant], BaseOp::Mult)?;
+
+    let neg_mask = not(config, region, &[relu_mask])?;
+
+    let quantized_alpha = quantize_tensor(
+        Tensor::from([alpha.0; 1].into_iter()),
+        *input_scale,
+        &crate::graph::Visibility::Fixed,
+    )?;
+
+    let alpha_tensor = create_constant_tensor(quantized_alpha[0], 1);
+
+    let scaled_neg_mask = pairwise(config, region, &[neg_mask, alpha_tensor], BaseOp::Mult)?;
+
+    let neg_part = pairwise(
+        config,
+        region,
+        &[values[0].clone(), scaled_neg_mask],
+        BaseOp::Mult,
+    )?;
+
+    pairwise(config, region, &[rescaled_positive, neg_part], BaseOp::Add)
 }
 
 fn multi_dim_axes_op<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(

src/circuit/ops/lookup.rs

@@ -15,101 +15,32 @@ use halo2curves::ff::PrimeField;
 /// An enum representing the operations that can be used to express more complex operations via accumulation
 #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Deserialize, Serialize)]
 pub enum LookupOp {
-    Div {
-        denom: utils::F32,
-    },
-    Cast {
-        scale: 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,
-    },
-    KroneckerDelta,
-    Pow {
-        scale: utils::F32,
-        a: utils::F32,
-    },
-    HardSwish {
-        scale: utils::F32,
-    },
+    Div { denom: utils::F32 },
+    Cast { scale: 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 },
+    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 },
+    Pow { scale: utils::F32, a: utils::F32 },
+    HardSwish { scale: utils::F32 },
 }
 
 impl LookupOp {
@@ -128,16 +59,8 @@ impl LookupOp {
             LookupOp::Round { scale } => format!("round_{}", scale),
             LookupOp::RoundHalfToEven { scale } => format!("round_half_to_even_{}", scale),
             LookupOp::Pow { scale, a } => format!("pow_{}_{}", scale, a),
-            LookupOp::KroneckerDelta => "kronecker_delta".into(),
-            LookupOp::Max { scale, a } => format!("max_{}_{}", scale, a),
-            LookupOp::Min { scale, a } => format!("min_{}_{}", scale, a),
             LookupOp::Div { denom } => format!("div_{}", denom),
             LookupOp::Cast { scale } => format!("cast_{}", scale),
-            LookupOp::Recip {
-                input_scale,
-                output_scale,
-            } => format!("recip_{}_{}", input_scale, output_scale),
-            LookupOp::LeakyReLU { slope: a } => format!("leaky_relu_{}", a),
             LookupOp::Sigmoid { scale } => format!("sigmoid_{}", scale),
             LookupOp::Sqrt { scale } => format!("sqrt_{}", scale),
             LookupOp::Rsqrt { scale } => format!("rsqrt_{}", scale),
@@ -183,32 +106,12 @@ impl LookupOp {
                 LookupOp::Pow { scale, a } => Ok::<_, TensorError>(
                     tensor::ops::nonlinearities::pow(&x, scale.0.into(), a.0.into()),
                 ),
-                LookupOp::KroneckerDelta => {
-                    Ok::<_, TensorError>(tensor::ops::nonlinearities::kronecker_delta(&x))
-                }
-                LookupOp::Max { scale, a } => Ok::<_, TensorError>(
-                    tensor::ops::nonlinearities::max(&x, scale.0.into(), a.0.into()),
-                ),
-                LookupOp::Min { scale, a } => Ok::<_, TensorError>(
-                    tensor::ops::nonlinearities::min(&x, scale.0.into(), a.0.into()),
-                ),
                 LookupOp::Div { denom } => Ok::<_, TensorError>(
                     tensor::ops::nonlinearities::const_div(&x, f32::from(*denom).into()),
                 ),
                 LookupOp::Cast { scale } => Ok::<_, TensorError>(
                     tensor::ops::nonlinearities::const_div(&x, f32::from(*scale).into()),
                 ),
-                LookupOp::Recip {
-                    input_scale,
-                    output_scale,
-                } => Ok::<_, TensorError>(tensor::ops::nonlinearities::recip(
-                    &x,
-                    input_scale.into(),
-                    output_scale.into(),
-                )),
-                LookupOp::LeakyReLU { slope: a } => {
-                    Ok::<_, TensorError>(tensor::ops::nonlinearities::leakyrelu(&x, a.0.into()))
-                }
                 LookupOp::Sigmoid { scale } => {
                     Ok::<_, TensorError>(tensor::ops::nonlinearities::sigmoid(&x, scale.into()))
                 }
@@ -288,20 +191,9 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Lookup
             LookupOp::Round { scale } => format!("ROUND(scale={})", scale),
             LookupOp::RoundHalfToEven { scale } => format!("ROUND_HALF_TO_EVEN(scale={})", scale),
             LookupOp::Pow { a, scale } => format!("POW(scale={}, exponent={})", scale, a),
-            LookupOp::KroneckerDelta => "K_DELTA".into(),
-            LookupOp::Max { scale, a } => format!("MAX(scale={}, a={})", scale, a),
-            LookupOp::Min { scale, a } => format!("MIN(scale={}, a={})", scale, a),
-            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),
-            LookupOp::LeakyReLU { slope: a } => format!("L_RELU(slope={})", a),
             LookupOp::Sigmoid { scale } => format!("SIGMOID(scale={})", scale),
             LookupOp::Sqrt { scale } => format!("SQRT(scale={})", scale),
             LookupOp::Erf { scale } => format!("ERF(scale={})", scale),
@@ -344,8 +236,6 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Lookup
                 let in_scale = inputs_scale[0];
                 in_scale + multiplier_to_scale(1. / scale.0 as f64)
             }
-            LookupOp::Recip { output_scale, .. } => multiplier_to_scale(output_scale.into()),
-            LookupOp::KroneckerDelta => 0,
             _ => inputs_scale[0],
         };
         Ok(scale)

src/circuit/ops/poly.rs

@@ -1,5 +1,8 @@
 use crate::{
-    circuit::layouts,
+    circuit::{
+        layouts,
+        utils::{self, F32},
+    },
     tensor::{self, Tensor, TensorError},
 };
 
@@ -9,9 +12,12 @@ use super::{base::BaseOp, *};
 /// An enum representing the operations that can be expressed as arithmetic (non lookup) operations.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub enum PolyOp {
-    ReLU,
     Abs,
     Sign,
+    LeakyReLU {
+        slope: utils::F32,
+        scale: i32,
+    },
     GatherElements {
         dim: usize,
         constant_idx: Option<Tensor<usize>>,
@@ -112,9 +118,9 @@ impl<
 
     fn as_string(&self) -> String {
         match &self {
+            PolyOp::LeakyReLU { slope: a, .. } => format!("LEAKYRELU (slope={})", a),
             PolyOp::Abs => "ABS".to_string(),
             PolyOp::Sign => "SIGN".to_string(),
-            PolyOp::ReLU => "RELU".to_string(),
             PolyOp::GatherElements { dim, constant_idx } => format!(
                 "GATHERELEMENTS (dim={}, constant_idx{})",
                 dim,
@@ -198,7 +204,9 @@ impl<
         Ok(Some(match self {
             PolyOp::Abs => layouts::abs(config, region, values[..].try_into()?)?,
             PolyOp::Sign => layouts::sign(config, region, values[..].try_into()?)?,
-            PolyOp::ReLU => layouts::relu(config, region, values[..].try_into()?)?,
+            PolyOp::LeakyReLU { slope, scale } => {
+                layouts::leaky_relu(config, region, values[..].try_into()?, slope, scale)?
+            }
             PolyOp::MultiBroadcastTo { shape } => {
                 layouts::expand(config, region, values[..].try_into()?, shape)?
             }
@@ -329,6 +337,12 @@ impl<
 
     fn out_scale(&self, in_scales: Vec<crate::Scale>) -> Result<crate::Scale, CircuitError> {
         let scale = match self {
+            // this corresponds to the relu operation
+            PolyOp::LeakyReLU {
+                slope: F32(0.0), ..
+            } => in_scales[0],
+            // this corresponds to the leaky relu operation with a slope which induces a change in scale
+            PolyOp::LeakyReLU { scale, .. } => in_scales[0] + *scale,
             PolyOp::MeanOfSquares { .. } => 2 * in_scales[0],
             PolyOp::Xor | PolyOp::Or | PolyOp::And | PolyOp::Not => 0,
             PolyOp::Iff => in_scales[1],

src/circuit/tests.rs

@@ -1379,7 +1379,10 @@ mod conv_relu_col_ultra_overflow {
                             .layout(
                                 &mut region,
                                 &[output.unwrap().unwrap()],
-                                Box::new(PolyOp::ReLU),
+                                Box::new(PolyOp::LeakyReLU {
+                                    slope: 0.0.into(),
+                                    scale: 1,
+                                }),
                             )
                             .unwrap();
                         Ok(())
@@ -2347,7 +2350,14 @@ mod matmul_relu {
                         .unwrap();
                     let _output = config
                         .base_config
-                        .layout(&mut region, &[output.unwrap()], Box::new(PolyOp::ReLU))
+                        .layout(
+                            &mut region,
+                            &[output.unwrap()],
+                            Box::new(PolyOp::LeakyReLU {
+                                slope: 0.0.into(),
+                                scale: 1,
+                            }),
+                        )
                         .unwrap();
                     Ok(())
                 },
@@ -2439,7 +2449,14 @@ mod relu {
                     |region| {
                         let mut region = RegionCtx::new(region, 0, 1, 2, 2);
                         Ok(config
-                            .layout(&mut region, &[self.input.clone()], Box::new(PolyOp::ReLU))
+                            .layout(
+                                &mut region,
+                                &[self.input.clone()],
+                                Box::new(PolyOp::LeakyReLU {
+                                    slope: 0.0.into(),
+                                    scale: 1,
+                                }),
+                            )
                             .unwrap())
                     },
                 )
@@ -2482,11 +2499,11 @@ mod lookup_ultra_overflow {
     use snark_verifier::system::halo2::transcript::evm::EvmTranscript;
 
     #[derive(Clone)]
-    struct ReLUCircuit<F: PrimeField + TensorType + PartialOrd> {
+    struct SigmoidCircuit<F: PrimeField + TensorType + PartialOrd> {
         pub input: ValTensor<F>,
     }
 
-    impl Circuit<F> for ReLUCircuit<F> {
+    impl Circuit<F> for SigmoidCircuit<F> {
         type Config = BaseConfig<F>;
         type FloorPlanner = SimpleFloorPlanner;
         type Params = TestParams;
@@ -2500,7 +2517,7 @@ mod lookup_ultra_overflow {
                 .map(|_| VarTensor::new_advice(cs, 4, 1, 3))
                 .collect::<Vec<_>>();
 
-            let nl = LookupOp::LeakyReLU { slope: 0.0.into() };
+            let nl = LookupOp::Sigmoid { scale: 1.0.into() };
 
             let mut config = BaseConfig::default();
 
@@ -2533,7 +2550,7 @@ mod lookup_ultra_overflow {
                             .layout(
                                 &mut region,
                                 &[self.input.clone()],
-                                Box::new(LookupOp::LeakyReLU { slope: 0.0.into() }),
+                                Box::new(LookupOp::Sigmoid { scale: 1.0.into() }),
                             )
                             .map_err(|_| Error::Synthesis)
                     },
@@ -2546,13 +2563,13 @@ mod lookup_ultra_overflow {
 
     #[test]
     #[ignore]
-    fn relucircuit() {
+    fn sigmoidcircuit() {
         // get some logs fam
         crate::logger::init_logger();
         // parameters
         let a = Tensor::from((0..4).map(|i| Value::known(F::from(i + 1))));
 
-        let circuit = ReLUCircuit::<F> {
+        let circuit = SigmoidCircuit::<F> {
             input: ValTensor::from(a),
         };
 
@@ -2562,7 +2579,7 @@ mod lookup_ultra_overflow {
 
         let pk = crate::pfsys::create_keys::<
             halo2_proofs::poly::kzg::commitment::KZGCommitmentScheme<halo2curves::bn256::Bn256>,
-            ReLUCircuit<F>,
+            SigmoidCircuit<F>,
         >(&circuit, &params, true)
         .unwrap();
 

src/graph/utilities.rs

@@ -763,81 +763,41 @@ pub fn new_op_from_onnx(
                 .map(|(i, _)| i)
                 .collect::<Vec<_>>();
 
-            if const_inputs.len() != 1 {
-                return Err(GraphError::OpMismatch(idx, "Max".to_string()));
-            }
-
-            let const_idx = const_inputs[0];
-            let boxed_op = inputs[const_idx].opkind();
-            let unit = if let Some(c) = extract_const_raw_values(boxed_op) {
-                if c.len() == 1 {
-                    c[0]
-                } else {
-                    return Err(GraphError::InvalidDims(idx, "max".to_string()));
-                }
-            } else {
-                return Err(GraphError::OpMismatch(idx, "Max".to_string()));
-            };
-
             if inputs.len() == 2 {
-                if let Some(node) = inputs.get_mut(const_idx) {
-                    node.decrement_use();
-                    deleted_indices.push(const_idx);
-                }
-                if unit == 0. {
-                    SupportedOp::Linear(PolyOp::ReLU)
+                if const_inputs.len() > 0 {
+                    let const_idx = const_inputs[0];
+                    let boxed_op = inputs[const_idx].opkind();
+                    let unit = if let Some(c) = extract_const_raw_values(boxed_op) {
+                        if c.len() == 1 {
+                            c[0]
+                        } else {
+                            return Err(GraphError::InvalidDims(idx, "max".to_string()));
+                        }
+                    } else {
+                        return Err(GraphError::OpMismatch(idx, "Max".to_string()));
+                    };
+                    if unit == 0. {
+                        if let Some(node) = inputs.get_mut(const_idx) {
+                            node.decrement_use();
+                            deleted_indices.push(const_idx);
+                        }
+                        SupportedOp::Linear(PolyOp::LeakyReLU {
+                            slope: 0.0.into(),
+                            scale: 1,
+                        })
+                    } else {
+                        SupportedOp::Hybrid(HybridOp::Max)
+                    }
                 } 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[non_const_idx].out_scales()[0]).into(),
-                        a: crate::circuit::utils::F32(unit),
-                    })
+                    SupportedOp::Hybrid(HybridOp::Max)
                 }
             } else {
                 return Err(GraphError::InvalidDims(idx, "max".to_string()));
             }
         }
         "Min" => {
-            // Extract the min value
-            // first find the input that is a constant
-            // and then extract the value
-            let const_inputs = inputs
-                .iter()
-                .enumerate()
-                .filter(|(_, n)| n.is_constant())
-                .map(|(i, _)| i)
-                .collect::<Vec<_>>();
-
-            if const_inputs.len() != 1 {
-                return Err(GraphError::OpMismatch(idx, "Min".to_string()));
-            }
-
-            let const_idx = const_inputs[0];
-            let boxed_op = inputs[const_idx].opkind();
-            let unit = if let Some(c) = extract_const_raw_values(boxed_op) {
-                if c.len() == 1 {
-                    c[0]
-                } else {
-                    return Err(GraphError::InvalidDims(idx, "min".to_string()));
-                }
-            } else {
-                return Err(GraphError::OpMismatch(idx, "Min".to_string()));
-            };
-
             if inputs.len() == 2 {
-                if let Some(node) = inputs.get_mut(const_idx) {
-                    node.decrement_use();
-                    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[non_const_idx].out_scales()[0]).into(),
-                    a: crate::circuit::utils::F32(unit),
-                })
+                SupportedOp::Hybrid(HybridOp::Min)
             } else {
                 return Err(GraphError::InvalidDims(idx, "min".to_string()));
             }
@@ -849,7 +809,6 @@ pub fn new_op_from_onnx(
             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: true,
             })
         }
 
@@ -864,8 +823,9 @@ pub fn new_op_from_onnx(
                 }
             };
 
-            SupportedOp::Nonlinear(LookupOp::LeakyReLU {
+            SupportedOp::Linear(PolyOp::LeakyReLU {
                 slope: crate::circuit::utils::F32(leaky_op.alpha),
+                scale: scales.params,
             })
         }
         "Scan" => {
@@ -1146,10 +1106,17 @@ pub fn new_op_from_onnx(
                 if c.raw_values.len() > 1 {
                     unimplemented!("only support scalar pow")
                 }
-                SupportedOp::Nonlinear(LookupOp::Pow {
-                    scale: scale_to_multiplier(inputs[0].out_scales()[0]).into(),
-                    a: crate::circuit::utils::F32(c.raw_values[0]),
-                })
+
+                let exponent = c.raw_values[0];
+
+                if exponent.fract() == 0.0 {
+                    SupportedOp::Linear(PolyOp::Pow(exponent as u32))
+                } else {
+                    SupportedOp::Nonlinear(LookupOp::Pow {
+                        scale: scale_to_multiplier(inputs[0].out_scales()[0]).into(),
+                        a: crate::circuit::utils::F32(exponent),
+                    })
+                }
             } else {
                 unimplemented!("only support constant pow for now")
             }

src/tensor/ops.rs

@@ -1553,35 +1553,6 @@ pub mod nonlinearities {
         .unwrap()
     }
 
-    /// Applies Kronecker delta to a tensor of integers.
-    /// # Arguments
-    /// * `a` - Tensor
-    /// # Examples
-    /// ```
-    /// use ezkl::tensor::Tensor;
-    /// use ezkl::fieldutils::IntegerRep;
-    /// use ezkl::tensor::ops::nonlinearities::kronecker_delta;
-    /// let x = Tensor::<IntegerRep>::new(
-    ///    Some(&[2, 15, 2, 1, 1, 0]),
-    ///  &[2, 3],
-    /// ).unwrap();
-    /// let result = kronecker_delta(&x);
-    /// let expected = Tensor::<IntegerRep>::new(Some(&[0, 0, 0, 0, 0, 1]), &[2, 3]).unwrap();
-    /// assert_eq!(result, expected);
-    /// ```
-    pub fn kronecker_delta<T: TensorType + std::cmp::PartialEq + Send + Sync>(
-        a: &Tensor<T>,
-    ) -> Tensor<T> {
-        a.par_enum_map(|_, a_i| {
-            if a_i == T::zero().unwrap() {
-                Ok::<_, TensorError>(T::one().unwrap())
-            } else {
-                Ok::<_, TensorError>(T::zero().unwrap())
-            }
-        })
-        .unwrap()
-    }
-
     /// Elementwise applies sigmoid to a tensor of integers.
     /// # Arguments
     ///

tests/integration_tests.rs

@@ -205,7 +205,7 @@ mod native_tests {
         "1l_tiny_div",
     ];
 
-    const TESTS: [&str; 94] = [
+    const TESTS: [&str; 95] = [
         "1l_mlp", //0
         "1l_slice",
         "1l_concat",
@@ -304,6 +304,7 @@ mod native_tests {
         "lstm_large",  // 91
         "lstm_medium", // 92
         "lenet_5",     // 93
+        "rsqrt",       // 94
     ];
 
     const WASM_TESTS: [&str; 46] = [
@@ -542,7 +543,7 @@ mod native_tests {
             }
         });
 
-            seq!(N in 0..=93 {
+            seq!(N in 0..=94 {
 
             #(#[test_case(TESTS[N])])*
             #[ignore]
@@ -1118,7 +1119,7 @@ mod native_tests {
 
             });
 
-            seq!(N in 0..=93 {
+            seq!(N in 0..4 {
                 #(#[test_case(TESTS[N])])*
                 fn kzg_evm_prove_and_verify_reusable_verifier_(test: &str) {
                     crate::native_tests::init_binary();

tests/py_integration_tests.rs

@@ -124,41 +124,40 @@ mod py_tests {
     }
 
     const TESTS: [&str; 34] = [
-        "ezkl_demo_batch.ipynb",
-        "proof_splitting.ipynb", // 0
-        "variance.ipynb",
-        "mnist_gan.ipynb",
-        // "mnist_vae.ipynb",
-        "keras_simple_demo.ipynb",
-        "mnist_gan_proof_splitting.ipynb", // 4
-        "hashed_vis.ipynb",                // 5
-        "simple_demo_all_public.ipynb",
-        "data_attest.ipynb",
-        "little_transformer.ipynb",
-        "simple_demo_aggregated_proofs.ipynb",
-        "ezkl_demo.ipynb", // 10
-        "lstm.ipynb",
-        "set_membership.ipynb", // 12
-        "decision_tree.ipynb",
-        "random_forest.ipynb",
-        "gradient_boosted_trees.ipynb", // 15
-        "xgboost.ipynb",
-        "lightgbm.ipynb",
-        "svm.ipynb",
-        "simple_demo_public_input_output.ipynb",
-        "simple_demo_public_network_output.ipynb", // 20
-        "gcn.ipynb",
-        "linear_regression.ipynb",
-        "stacked_regression.ipynb",
-        "data_attest_hashed.ipynb",
-        "kzg_vis.ipynb", // 25
-        "kmeans.ipynb",
-        "solvency.ipynb",
-        "sklearn_mlp.ipynb",
-        "generalized_inverse.ipynb",
-        "mnist_classifier.ipynb", // 30
-        "world_rotation.ipynb",
-        "logistic_regression.ipynb",
+        "ezkl_demo_batch.ipynb",                   // 0
+        "proof_splitting.ipynb",                   // 1
+        "variance.ipynb",                          // 2
+        "mnist_gan.ipynb",                         // 3
+        "keras_simple_demo.ipynb",                 // 4
+        "mnist_gan_proof_splitting.ipynb",         // 5
+        "hashed_vis.ipynb",                        // 6
+        "simple_demo_all_public.ipynb",            // 7
+        "data_attest.ipynb",                       // 8
+        "little_transformer.ipynb",                // 9
+        "simple_demo_aggregated_proofs.ipynb",     // 10
+        "ezkl_demo.ipynb",                         // 11
+        "lstm.ipynb",                              // 12
+        "set_membership.ipynb",                    // 13
+        "decision_tree.ipynb",                     // 14
+        "random_forest.ipynb",                     // 15
+        "gradient_boosted_trees.ipynb",            // 16
+        "xgboost.ipynb",                           // 17
+        "lightgbm.ipynb",                          // 18
+        "svm.ipynb",                               // 19
+        "simple_demo_public_input_output.ipynb",   // 20
+        "simple_demo_public_network_output.ipynb", // 21
+        "gcn.ipynb",                               // 22
+        "linear_regression.ipynb",                 // 23
+        "stacked_regression.ipynb",                // 24
+        "data_attest_hashed.ipynb",                // 25
+        "kzg_vis.ipynb",                           // 26
+        "kmeans.ipynb",                            // 27
+        "solvency.ipynb",                          // 28
+        "sklearn_mlp.ipynb",                       // 29
+        "generalized_inverse.ipynb",               // 30
+        "mnist_classifier.ipynb",                  // 31
+        "world_rotation.ipynb",                    // 32
+        "logistic_regression.ipynb",               // 33
     ];
 
     macro_rules! test_func {