JSTprove 1.1.0__py3-none-macosx_11_0_arm64.whl → 1.2.0__py3-none-macosx_11_0_arm64.whl

{jstprove-1.1.0.dist-info → jstprove-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: JSTprove
-Version: 1.1.0
+Version: 1.2.0
 Summary: Zero-knowledge proofs of ML inference on ONNX models
 Author: Inference Labs Inc
 Requires-Python: >=3.10
@@ -45,7 +45,7 @@ Dynamic: license-file
 Zero-knowledge proofs of ML inference on **ONNX** models — powered by [Polyhedra Network’s **Expander**](https://github.com/PolyhedraZK/Expander) (GKR/sum-check prover) and [**Expander Compiler Collection (ECC)**](https://github.com/PolyhedraZK/ExpanderCompilerCollection).
 
 * 🎯 **You bring ONNX** → we quantize, compile to a circuit, generate a witness, prove, and verify — via a simple CLI.
-* ✅ Supported ops (current): **Conv2D**, **GEMM/MatMul (FC)**, **ReLU**, **MaxPool2D**, **Add**.
+* ✅ Supported ops (current): **Conv2D**, **GEMM/MatMul (FC)**, **ReLU**, **MaxPool2D**, **Add**, **Mul**, **Sub**, **BatchNorm**.
 * 🧰 CLI details: see **[docs/cli.md](docs/cli.md)**
 
 👉 Just want to see it in action? Jump to [Quickstart (LeNet demo)](#quickstart-lenet-demo).<br>
@@ -85,7 +85,7 @@ You provide an **ONNX** model and inputs; JSTprove handles **quantization**, **c
 ### High-level architecture
 
 * **Python pipeline:** Converts **ONNX → quantized ONNX**, prepares I/O, drives the Rust runner, exposes the **CLI**.
-* **Rust crate:** `rust/jstprove_circuits` implements layer circuits (Conv2D, ReLU, MaxPool2D, GEMM/FC) and a runner.
+* **Rust crate:** `rust/jstprove_circuits` implements layer circuits (Conv2D, ReLU, MaxPool2D, GEMM/FC, BatchNorm) and a runner.
 * **Circuit frontend:** [ECC](https://github.com/PolyhedraZK/ExpanderCompilerCollection) Rust API for arithmetic circuits.
 * **Prover backend:** [Expander](https://github.com/PolyhedraZK/Expander) (GKR/sum-check prover/verification).
 

{jstprove-1.1.0.dist-info → jstprove-1.2.0.dist-info}/RECORD

@@ -1,9 +1,9 @@
-jstprove-1.1.0.dist-info/licenses/LICENSE,sha256=UXQRcYRUH-PfN27n3P-FMaZFY6jr9jFPKcwT7CWbljw,1160
+jstprove-1.2.0.dist-info/licenses/LICENSE,sha256=UXQRcYRUH-PfN27n3P-FMaZFY6jr9jFPKcwT7CWbljw,1160
 python/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/core/__init__.py,sha256=RlfbqGAaUulKl44QGMCkkGJBQZ8R_AgC5bU5zS7BjnA,97
 python/core/binaries/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/core/binaries/expander-exec,sha256=C_1JcezdfLp9sFOQ2z3wp2gcq1k8zjIR09CxJKGGIuM,7095168
-python/core/binaries/onnx_generic_circuit_1-1-0,sha256=2YBhVx-neun-Dmx3ntyLq20qwsLrY9coOcU2bNLprZ0,3086160
+python/core/binaries/onnx_generic_circuit_1-2-0,sha256=vLWr1O5PePljr54ZJ32dgHcuawzauRzuZpz7cZxvwgc,3144592
 python/core/circuit_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/core/circuit_models/generic_onnx.py,sha256=P65UZkfVBTE6YhaQ951S6QoTHPuU5ntDt8QL5pXghvw,8787
 python/core/circuit_models/simple_circuit.py,sha256=igQrZtQyreyHc26iAgCyDb0TuD2bJAoumYhc1pYPDzQ,4682
@@ -15,25 +15,30 @@ python/core/model_processing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NM
 python/core/model_processing/errors.py,sha256=uh2YFjuuU5JM3anMtSTLAH-zjlNAKStmLDZqRUgBWS8,4611
 python/core/model_processing/converters/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/core/model_processing/converters/base.py,sha256=eG7iRDbDJJDTG2cCVgYlPlfkpmYPEnMzjGNK9wrA1m0,4303
-python/core/model_processing/converters/onnx_converter.py,sha256=BJc6rU3wLHI3imt8yzm8Cngri3KvcBSUbJ3Urw2PoEQ,44560
+python/core/model_processing/converters/onnx_converter.py,sha256=-eXdF6tfluFRxGgnQtJQ8R2309aYX-8z8HzMxk_Qv8I,44340
 python/core/model_processing/onnx_custom_ops/__init__.py,sha256=ofecV9pzpDJJl_r6inRw9JOKxtfK2rzzxWahAq9BKXE,475
+python/core/model_processing/onnx_custom_ops/batchnorm.py,sha256=8kg4iGGdt6B_fIJkpt4v5eNFpoHa4bjTB0NnCSmKFvE,1693
 python/core/model_processing/onnx_custom_ops/conv.py,sha256=6jJm3fcGWzcU4RjVgf179mPFCqsl4C3AR7bqQTffDgA,3464
 python/core/model_processing/onnx_custom_ops/custom_helpers.py,sha256=2WdnHw9NAoN_6wjIBoAQDyL6wEIlZOqo6ysCZp5DpZs,1844
 python/core/model_processing/onnx_custom_ops/gemm.py,sha256=bnEUXhqQCEcH4TIfbMTsCTtAlAlRzFvl4jj8g2QZFWU,2674
 python/core/model_processing/onnx_custom_ops/maxpool.py,sha256=Sd3BwqpGLSVU2iuAAIXAHdI3WO27Aa3g3r29HPiECvM,2319
+python/core/model_processing/onnx_custom_ops/mul.py,sha256=w6X1sl1HnzoUJx2Mm_LaoXGTpvtwXxr3zZDPySVHBcM,1888
 python/core/model_processing/onnx_custom_ops/onnx_helpers.py,sha256=utnJuc5sgb_z1LgxuY9y2cQbMpdEJ8xOOrcP8DhfDCM,5686
 python/core/model_processing/onnx_custom_ops/relu.py,sha256=pZsPXC_r0FPggURKDphh8P1IRXY0w4hH7ExBmYTlWjE,1202
 python/core/model_processing/onnx_quantizer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/core/model_processing/onnx_quantizer/exceptions.py,sha256=_YaXXEMbfD1P8N86L5YIz3uCilkuzlhv_2lU90T4FfA,5646
-python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py,sha256=POoDEBFzkr145P4INgAux2LQY2GdpsBtRpw_UuKVNhw,7679
+python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py,sha256=ncL0rK5hXZUvssmw20PZO1WyjYSyenem23B6QLUHlLY,9213
 python/core/model_processing/onnx_quantizer/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/core/model_processing/onnx_quantizer/layers/add.py,sha256=AGxzqMa0jABIEKOIgPqEAA7EpZtynQtnD9nxI2NHc0s,1409
-python/core/model_processing/onnx_quantizer/layers/base.py,sha256=LvyTvmR2w6jYSJiBvyFluaDgL_Voc6dZ00TTWi6V7Tc,17426
+python/core/model_processing/onnx_quantizer/layers/base.py,sha256=Vq6pwChw9eJMKYAJyA1C3wLycaBConkP9sNRInpWavo,19989
+python/core/model_processing/onnx_quantizer/layers/batchnorm.py,sha256=KSBDPHd52f5Qyf-cnIDFPmfzssaJgMPiTmpIWEdM41U,7718
 python/core/model_processing/onnx_quantizer/layers/constant.py,sha256=l1IvgvXkmFMiaBsym8wchPF-y1ZH-c5PmFUy92IXWok,3694
 python/core/model_processing/onnx_quantizer/layers/conv.py,sha256=TlUpCRO6PPqH7MPkIrEiEcVfzuiN1WMYEiNIjhYXtWM,4451
 python/core/model_processing/onnx_quantizer/layers/gemm.py,sha256=7fCUMv8OLVZ45a2lYjA2XNvcW3By7lSbX7zeForNK-0,3950
 python/core/model_processing/onnx_quantizer/layers/maxpool.py,sha256=PJ8hZPPBpfWV_RZdySl50-BU8TATjcg8Tg_mrAVS1Ic,4916
+python/core/model_processing/onnx_quantizer/layers/mul.py,sha256=qHsmnYPH-c5uiFeDCvV6e1xSgmIXJ64Sjvh0LYDYEqQ,1396
 python/core/model_processing/onnx_quantizer/layers/relu.py,sha256=d-5fyeKNLTgKKnqCwURpxkjl7QdbJQpuovtCFBM03FA,1685
+python/core/model_processing/onnx_quantizer/layers/sub.py,sha256=M7D98TZBNP9-2R9MX6mcpYlrWFxTiX9JCs3XNcg1U-Q,1409
 python/core/model_templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/core/model_templates/circuit_template.py,sha256=X8bA4AdmtQeb3ltU74GaWYfrOFhqs_DOpUqRMFXLAD8,2352
 python/core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -62,7 +67,7 @@ python/frontend/commands/bench/model.py,sha256=SaIWXAXZbWGbrNqEo5bs4NwgZfMOmmxaC
 python/frontend/commands/bench/sweep.py,sha256=rl-QBS9eXgQkuPJBhsU4CohfE1PdJvnM8NRhNU7ztQw,5279
 python/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/scripts/benchmark_runner.py,sha256=sjbqaLrdjt94AoyQXAxT4FhsN6aRu5idTRQ5uHmZOWM,28593
-python/scripts/gen_and_bench.py,sha256=9kcIj-K_nG-G194C68Uig-Yw-p3nYKESACIpWRflmts,16276
+python/scripts/gen_and_bench.py,sha256=V36x7djYmHlveAJgYzMlXwnmF0gAGO3-1mg9PWOmpj8,16249
 python/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/tests/test_cli.py,sha256=OiAyG3aBpukk0i5FFWbiKaF42wf-7By-UWDHNjwtsqo,27042
 python/tests/circuit_e2e_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -82,27 +87,30 @@ python/tests/onnx_quantizer_tests/testing_helper_functions.py,sha256=N0fQv2pYzUC
 python/tests/onnx_quantizer_tests/layers/__init__.py,sha256=xP-RmW6LfIANgK1s9Q0KZet2yvNr-3c6YIVLAAQqGUY,404
 python/tests/onnx_quantizer_tests/layers/add_config.py,sha256=T3tGddupDtrvLck2SL2yETDblNtv0aU7Tl7fNyZUhO4,4133
 python/tests/onnx_quantizer_tests/layers/base.py,sha256=uLCqhMcBA7zWiRSLRMNKKb4A9N27l-RUqSEEQ8SR3xI,9393
+python/tests/onnx_quantizer_tests/layers/batchnorm_config.py,sha256=P-sZuHAdEfNczcgTeLjqJnEbpqN3dKTsbqvY4-SBqiQ,8231
 python/tests/onnx_quantizer_tests/layers/constant_config.py,sha256=RdrKNMNZjI3Sk5o8WLNqmBUyYVJRWgtFbQ6oFWMwyQk,1193
 python/tests/onnx_quantizer_tests/layers/conv_config.py,sha256=H0ioW4H3ei5IK4tKhrA0ffThxJ4K5oO9jIs9A0T0VaM,6005
 python/tests/onnx_quantizer_tests/layers/factory.py,sha256=WLLEP9ECmSpTliSjhtdWOHcX1xOi6HM10S9Y4re1A74,4844
 python/tests/onnx_quantizer_tests/layers/flatten_config.py,sha256=Xln5Hh6gyeM5gGRCjLGvIL-u08NEs1tXSF32urCqPfE,2110
 python/tests/onnx_quantizer_tests/layers/gemm_config.py,sha256=t7nJY-Wnj6YUD821-jaWzgrQVPa6ytwER3hFMsvyY6Y,7294
 python/tests/onnx_quantizer_tests/layers/maxpool_config.py,sha256=XfTPk_ZQXEzaCjHHymSLVv2HS-PKH1rS9IuyyoEtM78,3176
+python/tests/onnx_quantizer_tests/layers/mul_config.py,sha256=_Oy4b97ORxFlF3w0BmJ94hNA968HQx2AvwYiASrGPxw,4135
 python/tests/onnx_quantizer_tests/layers/relu_config.py,sha256=_aHuddDApLUBOa0FiR9h4fNfmMSnH5r4JzOMLW0KaTk,2197
 python/tests/onnx_quantizer_tests/layers/reshape_config.py,sha256=fZchSqIAy76m7j97wVC_UI6slSpv8nbwukhkbGR2sRE,2203
+python/tests/onnx_quantizer_tests/layers/sub_config.py,sha256=IxF18mG9kjlEiKYSNG912CEcBxOFGxIWoRAwjvBXiRo,4133
 python/tests/onnx_quantizer_tests/layers_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/tests/onnx_quantizer_tests/layers_tests/base_test.py,sha256=UgbcT97tgcuTtO1pOADpww9bz_JElKiI2mxLJYKyF1k,2992
 python/tests/onnx_quantizer_tests/layers_tests/test_check_model.py,sha256=Vxn4LEWHZeGa_vS1-7ptFqSSBb0D-3BG-ETocP4pvsI,3651
 python/tests/onnx_quantizer_tests/layers_tests/test_e2e.py,sha256=40779aaHgdryVwLlIO18F1d7uSLSXdJUG5Uj_5-xD4U,6712
 python/tests/onnx_quantizer_tests/layers_tests/test_error_cases.py,sha256=t5c_zqO4Ex3HIFWcykX4PTftdKN7UWnEOF5blShL0Ik,1881
 python/tests/onnx_quantizer_tests/layers_tests/test_integration.py,sha256=Mq1-PBKR3756i9VrFOP5DY3GkRE32D6Hjd1fK9wZdVk,7228
-python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py,sha256=zclzXxtgA5BEmNwSf_aNbJgbsArMXn5WDdlxiMR2-aM,9255
+python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py,sha256=bVdMDkIq0gdHNLTFrWRdrCgCAG03rEF8aCRU-t4b4Kg,9391
 python/tests/onnx_quantizer_tests/layers_tests/test_scalability.py,sha256=RfnIIiYbgPbU3620H6MPvSxE3MNR2G1yPELwdWV3mK4,4107
 python/tests/onnx_quantizer_tests/layers_tests/test_validation.py,sha256=jz-WtIEP-jjUklOOAnznwPUXbf07U2PAMGrhzMWP0JU,1371
 python/tests/utils_testing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/tests/utils_testing/test_helper_functions.py,sha256=xmeGQieh4LE9U-CDKBlHhSWqH0cAmmDU3qXNbDkkvms,27192
-jstprove-1.1.0.dist-info/METADATA,sha256=3gdOLaD4eYGawv4SuvofjuzBW-y564J4gpNPXHFNY1A,14056
-jstprove-1.1.0.dist-info/WHEEL,sha256=jc2C2uw104ioj1TL9cE0YO67_kdAwX4W8JgYPomxr5M,105
-jstprove-1.1.0.dist-info/entry_points.txt,sha256=nGcTSO-4q08gPl1IoWdrPaiY7IbO7XvmXKkd34dYHc8,49
-jstprove-1.1.0.dist-info/top_level.txt,sha256=J-z0poNcsv31IHB413--iOY8LoHBKiTHeybHX3abokI,7
-jstprove-1.1.0.dist-info/RECORD,,
+jstprove-1.2.0.dist-info/METADATA,sha256=UVxR8iFm2kjvrvh1t4hEaCn0n4ZCYE2fcurGeCRmRCk,14100
+jstprove-1.2.0.dist-info/WHEEL,sha256=jc2C2uw104ioj1TL9cE0YO67_kdAwX4W8JgYPomxr5M,105
+jstprove-1.2.0.dist-info/entry_points.txt,sha256=nGcTSO-4q08gPl1IoWdrPaiY7IbO7XvmXKkd34dYHc8,49
+jstprove-1.2.0.dist-info/top_level.txt,sha256=J-z0poNcsv31IHB413--iOY8LoHBKiTHeybHX3abokI,7
+jstprove-1.2.0.dist-info/RECORD,,

python/core/model_processing/converters/onnx_converter.py

@@ -247,6 +247,7 @@ class ONNXConverter(ModelConverter):
 
     def analyze_layers(
         self: ONNXConverter,
+        model: onnx.ModelProto,
         output_name_to_shape: dict[str, list[int]] | None = None,
     ) -> tuple[list[ONNXLayer], list[ONNXLayer]]:
         """Analyze the onnx model graph into
@@ -268,29 +269,29 @@ class ONNXConverter(ModelConverter):
             id_count = 0
             # Apply shape inference on the model
             if not output_name_to_shape:
-                inferred_model = shape_inference.infer_shapes(self.model)
+                inferred_model = shape_inference.infer_shapes(model)
                 self._onnx_check_model_safely(inferred_model)
 
                 output_name_to_shape = extract_shape_dict(inferred_model)
             domain_to_version = {
-                opset.domain: opset.version for opset in self.model.opset_import
+                opset.domain: opset.version for opset in model.opset_import
             }
 
             id_count = 0
             architecture = self.get_model_architecture(
-                self.model,
+                model,
                 output_name_to_shape,
                 domain_to_version,
             )
             w_and_b = self.get_model_w_and_b(
-                self.model,
+                model,
                 output_name_to_shape,
                 id_count,
                 domain_to_version,
             )
         except InvalidModelError:
             raise
-        except (ValueError, TypeError, RuntimeError, OSError, onnx.ONNXException) as e:
+        except (ValueError, TypeError, RuntimeError, OSError) as e:
             raise LayerAnalysisError(model_type=self.model_type, reason=str(e)) from e
         except Exception as e:
             raise LayerAnalysisError(model_type=self.model_type, reason=str(e)) from e
@@ -557,6 +558,7 @@ class ONNXConverter(ModelConverter):
         output_shapes = {
             out_name: output_name_to_shape.get(out_name, []) for out_name in outputs
         }
+
         return ONNXLayer(
             id=layer_id,
             name=name,
@@ -605,6 +607,7 @@ class ONNXConverter(ModelConverter):
             np_data = onnx.numpy_helper.to_array(node, constant_dtype)
         except (ValueError, TypeError, onnx.ONNXException, Exception) as e:
             raise SerializationError(
+                model_type=self.model_type,
                 tensor_name=node.name,
                 reason=f"Failed to convert tensor: {e!s}",
             ) from e
@@ -1040,38 +1043,36 @@ class ONNXConverter(ModelConverter):
                   ``rescale_config``.
         """
         inferred_model = shape_inference.infer_shapes(self.model)
-
-        scaling = BaseOpQuantizer.get_scaling(
-            scale_base=getattr(self, "scale_base", 2),
-            scale_exponent=(getattr(self, "scale_exponent", 18)),
-        )
+        scale_base = getattr(self, "scale_base", 2)
+        scale_exponent = getattr(self, "scale_exponent", 18)
 
         # Check the model and print Y"s shape information
         self._onnx_check_model_safely(inferred_model)
         output_name_to_shape = extract_shape_dict(inferred_model)
-        (architecture, w_and_b) = self.analyze_layers(output_name_to_shape)
-        for w in w_and_b:
+        scaled_and_transformed_model = self.op_quantizer.apply_pre_analysis_transforms(
+            inferred_model,
+            scale_exponent=scale_exponent,
+            scale_base=scale_base,
+        )
+        # Get layers in correct format
+        (architecture, w_and_b) = self.analyze_layers(
+            scaled_and_transformed_model,
+            output_name_to_shape,
+        )
+
+        def _convert_tensor_to_int_list(w: ONNXLayer) -> list:
             try:
-                w_and_b_array = np.asarray(w.tensor)
-            except (ValueError, TypeError, Exception) as e:
+                arr = np.asarray(w.tensor).astype(np.int64)
+                return arr.tolist()
+            except Exception as e:
                 raise SerializationError(
                     tensor_name=getattr(w, "name", None),
+                    model_type=self.model_type,
                     reason=f"cannot convert to ndarray: {e}",
                 ) from e
 
-            try:
-                # TODO @jsgold-1: We need a better way to distinguish bias tensors from weight tensors # noqa: FIX002, TD003,E501
-                if "bias" in w.name:
-                    w_and_b_scaled = w_and_b_array * scaling * scaling
-                else:
-                    w_and_b_scaled = w_and_b_array * scaling
-                w_and_b_out = w_and_b_scaled.astype(np.int64).tolist()
-                w.tensor = w_and_b_out
-            except (ValueError, TypeError, OverflowError, Exception) as e:
-                raise SerializationError(
-                    tensor_name=getattr(w, "name", None),
-                    reason=str(e),
-                ) from e
+        for w in w_and_b:
+            w.tensor = _convert_tensor_to_int_list(w)
 
         inputs = []
         outputs = []

python/core/model_processing/onnx_custom_ops/batchnorm.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import numpy as np
+from onnxruntime_extensions import PyCustomOpDef, onnx_op
+
+from .custom_helpers import rescaling
+
+
+@onnx_op(
+    op_type="Int64BatchNorm",
+    domain="ai.onnx.contrib",
+    inputs=[
+        PyCustomOpDef.dt_int64,  # X (int64)
+        PyCustomOpDef.dt_int64,  # mul (int64 scaled multiplier)
+        PyCustomOpDef.dt_int64,  # add (int64 scaled adder)
+        PyCustomOpDef.dt_int64,  # scaling_factor
+    ],
+    outputs=[PyCustomOpDef.dt_int64],
+    attrs={"rescale": PyCustomOpDef.dt_int64},
+)
+def int64_batchnorm(
+    x: np.ndarray,
+    mul: np.ndarray,
+    add: np.ndarray,
+    scaling_factor: np.ndarray | None = None,
+    rescale: int | None = None,
+) -> np.ndarray:
+    """
+    Int64 BatchNorm (folded into affine transform).
+
+    Computes:
+        Y = X * mul + add
+    where mul/add are already scaled to int64.
+
+    Parameters
+    ----------
+    x : Input int64 tensor
+    mul : Per-channel int64 scale multipliers
+    add : Per-channel int64 bias terms
+    scaling_factor: factor to rescale
+    rescale : Optional flag to apply post-scaling
+
+    Returns
+    -------
+    numpy.ndarray (int64)
+    """
+    try:
+        # Broadcasting shapes must match batchnorm layout: NCHW
+        # Typically mul/add have shape [C]
+        dims_x = len(x.shape)
+        dim_ones = (1,) * (dims_x - 2)
+        mul = mul.reshape(-1, *dim_ones)
+        add = add.reshape(-1, *dim_ones)
+
+        y = x * mul + add
+
+        if rescale is not None:
+            y = rescaling(scaling_factor, rescale, y)
+
+        return y.astype(np.int64)
+
+    except Exception as e:
+        msg = f"Int64BatchNorm failed: {e}"
+        raise RuntimeError(msg) from e

python/core/model_processing/onnx_custom_ops/mul.py

@@ -0,0 +1,66 @@
+import numpy as np
+from onnxruntime_extensions import PyCustomOpDef, onnx_op
+
+from .custom_helpers import rescaling
+
+
+@onnx_op(
+    op_type="Int64Mul",
+    domain="ai.onnx.contrib",
+    inputs=[
+        PyCustomOpDef.dt_int64,
+        PyCustomOpDef.dt_int64,
+        PyCustomOpDef.dt_int64,  # Scalar
+    ],
+    outputs=[PyCustomOpDef.dt_int64],
+    attrs={
+        "rescale": PyCustomOpDef.dt_int64,
+    },
+)
+def int64_mul(
+    a: np.ndarray,
+    b: np.ndarray,
+    scaling_factor: np.ndarray | None = None,
+    rescale: int | None = None,
+) -> np.ndarray:
+    """
+    Performs a Mul (hadamard product) operation on int64 input tensors.
+
+    This function is registered as a custom ONNX operator via onnxruntime_extensions
+    and is used in the JSTprove quantized inference pipeline.
+    It applies Mul with the rescaling the outputs back to the original scale.
+
+    Parameters
+    ----------
+    a : np.ndarray
+        First input tensor with dtype int64.
+    b : np.ndarray
+        Second input tensor with dtype int64.
+    scaling_factor : Scaling factor for rescaling the output.
+        Optional scalar tensor for rescaling when rescale=1.
+    rescale : int, optional
+        Whether to apply rescaling (0=no, 1=yes).
+
+    Returns
+    -------
+    numpy.ndarray
+        Mul tensor with dtype int64.
+
+    Notes
+    -----
+    - This op is part of the `ai.onnx.contrib` custom domain.
+    - ONNX Runtime Extensions is required to register this op.
+
+    References
+    ----------
+    For more information on the Mul operation, please refer to the
+    ONNX standard Mul operator documentation:
+    https://onnx.ai/onnx/operators/onnx__Mul.html
+    """
+    try:
+        result = a * b
+        result = rescaling(scaling_factor, rescale, result)
+        return result.astype(np.int64)
+    except Exception as e:
+        msg = f"Int64Mul failed: {e}"
+        raise RuntimeError(msg) from e

python/core/model_processing/onnx_quantizer/layers/base.py

@@ -479,6 +479,73 @@ class QuantizerBase:
         nodes.append(quantized_node)
         return nodes
 
+    def pre_analysis_transform(
+        self: QuantizerBase,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        initializer_map: dict[str, onnx.TensorProto],
+        scale_base: int,
+        scale_exponent: int,
+    ) -> None:
+        """
+        pre_analysis_transform aims to transform the given layer along the
+        same lines as it would be transformed for the quantized model, but
+        for the weights and biases file instead, to be sent to the backend
+
+        Default pre-analysis behavior:
+
+        - If the subclass uses weights/bias (`USE_WB=True`), apply the SAME
+        scaling rules as quantization, but directly mutate the initializers.
+
+        - Subclasses can override this to implement more complex rewrites
+        (e.g., BatchNorm → Mul/Add).
+
+        Args:
+            node (onnx.NodeProto): Node to transform.
+            graph (onnx.GraphProto): Rest of the Onnx graph for initializers.
+            initializer_map (dict[str, onnx.TensorProto]): The initializer map.
+
+            scale_base (int): Scaling base.
+            scale_exponent (int): Scaling exponent.
+
+        NOTE
+         - The resulting model will not make accurate prediction and should be
+         used solely for analysis and keeping track of w_and_b
+        """
+        # If subclass does not want auto-scaling, do nothing
+        if not getattr(self, "USE_WB", False):
+            return
+
+        # Each quantizer defines which inputs to scale (Weight:1x, Bias:2x etc.)
+        scale_plan = getattr(self, "SCALE_PLAN", {})
+
+        # Perform the same scaling as quantization, but directly modify initializers
+        for input_idx, scale_mult in scale_plan.items():
+            if input_idx >= len(node.input):
+                continue
+
+            name = node.input[input_idx]
+            if name not in initializer_map:
+                continue  # optional input missing
+
+            tensor = initializer_map[name]
+            arr = numpy_helper.to_array(tensor).astype(np.float64)
+
+            scale = scale_base ** (scale_exponent * scale_mult)
+            new_arr = arr * scale
+
+            # Replace initializer directly
+            new_tensor = numpy_helper.from_array(new_arr, name=tensor.name)
+
+            # Modify graph initializer in place
+            for j in range(len(graph.initializer)):
+                if graph.initializer[j].name == tensor.name:
+                    del graph.initializer[j]
+                    break
+            graph.initializer.append(new_tensor)
+
+            initializer_map[tensor.name] = new_tensor
+
 
 class PassthroughQuantizer(BaseOpQuantizer):
     """

python/core/model_processing/onnx_quantizer/layers/batchnorm.py

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+from python.core.circuits.errors import CircuitConfigurationError
+
+if TYPE_CHECKING:
+    import onnx
+
+import numpy as np
+from onnx import helper, numpy_helper
+
+from python.core.model_processing.onnx_custom_ops.onnx_helpers import extract_attributes
+from python.core.model_processing.onnx_quantizer.exceptions import InvalidParamError
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    QuantizerBase,
+    ScaleConfig,
+)
+
+
+class QuantizeBatchnorm(QuantizerBase):
+    OP_TYPE = "Int64BatchNorm"
+    USE_WB = True
+    USE_SCALING = False
+    SCALE_PLAN: ClassVar = {}
+
+
+class BatchnormQuantizer(BaseOpQuantizer, QuantizeBatchnorm):
+    """
+    Quantizer for ONNX Batchnorm layers.
+
+    - Uses standard ONNX Batchnorm layer in standard domain, and
+      makes relevant additional changes to the graph.
+    """
+
+    def __init__(
+        self: BatchnormQuantizer,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        # Only replace if caller provided something
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
+
+    def _compute_mul_add(
+        self: BatchnormQuantizer,
+        initializer_map: dict[str, onnx.TensorProto],
+        node: onnx.NodeProto,
+        scale_base: int,
+        scale_exponent: int,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Compute the 'mul' and 'add' tensors for BatchNorm folding.
+        """
+        self._validate_inputs(node=node)
+        # ONNX BatchNorm inputs: [X, scale, bias, mean, var]
+        scale_factor = scale_base**scale_exponent
+        scale = numpy_helper.to_array(initializer_map[node.input[1]]).astype(np.float32)
+        bias = numpy_helper.to_array(initializer_map[node.input[2]]).astype(np.float32)
+        mean = numpy_helper.to_array(initializer_map[node.input[3]]).astype(np.float32)
+        var = numpy_helper.to_array(initializer_map[node.input[4]]).astype(np.float32)
+
+        # Find epsilon attribute
+        epsilon_attr = next((a for a in node.attribute if a.name == "epsilon"), None)
+        epsilon = float(epsilon_attr.f) if epsilon_attr else 1e-5
+
+        mul = scale / np.sqrt(var + epsilon)
+        add = bias - mean * mul
+        scaled_add = add * (scale_factor**2)
+        scaled_mul = scale_factor * mul
+        return scaled_mul, scaled_add
+
+    def pre_analysis_transform(
+        self: BatchnormQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        initializer_map: dict[str, onnx.TensorProto],
+        scale_base: int,
+        scale_exponent: int,
+    ) -> None:
+        # Compute linearized BN tensors
+        mul, add = self._compute_mul_add(
+            initializer_map,
+            node,
+            scale_base=scale_base,
+            scale_exponent=scale_exponent,
+        )
+
+        # Name base
+        node_name = node.name if node.name else node.input[0]
+        mul_name = f"{node_name}_mul"
+        add_name = f"{node_name}_add"
+
+        # Create ONNX tensors
+        mul_tensor = numpy_helper.from_array(mul.astype(np.int64), name=mul_name)
+        add_tensor = numpy_helper.from_array(add.astype(np.int64), name=add_name)
+
+        # Insert them into the graph
+        graph.initializer.extend([mul_tensor, add_tensor])
+        initializer_map[mul_name] = mul_tensor
+        initializer_map[add_name] = add_tensor
+        self.new_initializers.extend([mul_tensor, add_tensor])
+
+        node.input[:] = [node.input[0], mul_name, add_name]
+
+        del node.attribute[:]
+
+    def quantize(
+        self,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        _ = graph
+
+        nodes: list[onnx.NodeProto] = []
+
+        # 1. Compute unscaled float mul/add coefficients
+        mul, add = self._compute_mul_add(
+            initializer_map,
+            node,
+            scale_base=1,
+            scale_exponent=1,
+        )
+
+        node_name = node.name if node.name else node.input[0]
+        mul_name = f"{node_name}_mul"
+        add_name = f"{node_name}_add"
+
+        # 2. Store unscaled mul and add initializers (as floats)
+        scale_value = self.get_scaling(scale_config.base, scale_config.exponent)
+        scale_name = f"{node.name}_int_scaler"
+        scale_tensor = numpy_helper.from_array(
+            np.array([scale_value], dtype=np.int64),
+            name=scale_name,
+        )
+        self.new_initializers.append(scale_tensor)
+
+        mul_tensor = numpy_helper.from_array(mul.astype(np.float32), name=mul_name)
+        add_tensor = numpy_helper.from_array(add.astype(np.float32), name=add_name)
+
+        initializer_map[mul_name] = mul_tensor
+        initializer_map[add_name] = add_tensor
+
+        # 3. Insert scale and cast for mul_tensor
+        scaled_mul_name, mul_scale_node, mul_cast_node = self.insert_scale_node(
+            tensor=mul_tensor,
+            scale_base=scale_config.base,
+            scale_exponent=scale_config.exponent,
+        )
+
+        # 4. Insert scale and cast for add_tensor
+        scaled_add_name, add_scale_node, add_cast_node = self.insert_scale_node(
+            tensor=add_tensor,
+            scale_base=scale_config.base,
+            scale_exponent=scale_config.exponent * 2,
+        )
+        # Note, order is important here
+        nodes.extend(
+            [
+                mul_scale_node,
+                mul_cast_node,
+                add_scale_node,
+                add_cast_node,
+            ],
+        )
+
+        # 5. Build final Int64BatchNorm node
+        attrs = extract_attributes(node)
+        for k, v in getattr(self, "DEFAULT_ATTRS", {}).items():
+            attrs.setdefault(k, v)
+        attrs["rescale"] = 1
+
+        quant_node = helper.make_node(
+            self.OP_TYPE,  # Should be "Int64BatchNorm"
+            inputs=[
+                node.input[0],  # original X
+                scaled_mul_name,  # scaled mul
+                scaled_add_name,  # scaled add
+                scale_name,  # scaling factor
+            ],
+            outputs=node.output,
+            name=node.name,
+            domain=self.DOMAIN,
+            **attrs,
+        )
+
+        nodes.append(quant_node)
+        return nodes
+
+    def check_supported(
+        self: BatchnormQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        """
+        For our current implementation, all batchnorm inputs
+        (scale, variance, mean, etc.)
+        must be initializers to the circuit and not inputs from earlier in the graph.
+        """
+
+        if initializer_map is None:
+            msg = "initializer_map is required for BatchNorm support check"
+            raise CircuitConfigurationError(node.name, node.op_type, msg)
+
+        self._validate_inputs(node=node)
+
+        # First, check to make sure that each of the batchnorm inputs are initializers
+        initializer_inputs = node.input[1:]
+        if not all(i in initializer_map for i in initializer_inputs):
+            msg = "Unsupported BatchNorm with normalization inputs not in initializers"
+            raise InvalidParamError(node.name, node.op_type, msg)
+
+    def _validate_inputs(self, node: onnx.NodeProto) -> None:
+        """Validate BatchNorm has required inputs in initializer_map."""
+        num_inputs = 5
+        if len(node.input) < num_inputs:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"BatchNorm requires 5 inputs, got {len(node.input)}",
+            )

python/core/model_processing/onnx_quantizer/layers/mul.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import onnx
+
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    QuantizerBase,
+    ScaleConfig,
+)
+
+
+class QuantizeMul(QuantizerBase):
+    OP_TYPE = "Int64Mul"
+    USE_WB = True
+    USE_SCALING = True
+    SCALE_PLAN: ClassVar = {0: 1, 1: 1}
+
+
+class MulQuantizer(BaseOpQuantizer, QuantizeMul):
+    """
+    Quantizer for ONNX Mul layers.
+
+    - Uses custom Mul layer to incorporate rescaling, and
+      makes relevant additional changes to the graph.
+    """
+
+    def __init__(
+        self: MulQuantizer,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        # Only replace if caller provided something
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
+
+    def quantize(
+        self: MulQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        return QuantizeMul.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self: MulQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        pass

python/core/model_processing/onnx_quantizer/layers/sub.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import onnx
+
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    QuantizerBase,
+    ScaleConfig,
+)
+
+
+class QuantizeSub(QuantizerBase):
+    OP_TYPE = "Sub"
+    DOMAIN = ""
+    USE_WB = True
+    USE_SCALING = False
+    SCALE_PLAN: ClassVar = {0: 1, 1: 1}
+
+
+class SubQuantizer(BaseOpQuantizer, QuantizeSub):
+    """
+    Quantizer for ONNX Sub layers.
+
+    - Uses standard ONNX Sub layer in standard domain, and
+      makes relevant additional changes to the graph.
+    """
+
+    def __init__(
+        self: SubQuantizer,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        # Only replace if caller provided something
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
+
+    def quantize(
+        self: SubQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        return QuantizeSub.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self: SubQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        pass

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py

@@ -17,13 +17,18 @@ from python.core.model_processing.onnx_quantizer.layers.base import (
     PassthroughQuantizer,
     ScaleConfig,
 )
+from python.core.model_processing.onnx_quantizer.layers.batchnorm import (
+    BatchnormQuantizer,
+)
 from python.core.model_processing.onnx_quantizer.layers.constant import (
     ConstantQuantizer,
 )
 from python.core.model_processing.onnx_quantizer.layers.conv import ConvQuantizer
 from python.core.model_processing.onnx_quantizer.layers.gemm import GemmQuantizer
 from python.core.model_processing.onnx_quantizer.layers.maxpool import MaxpoolQuantizer
+from python.core.model_processing.onnx_quantizer.layers.mul import MulQuantizer
 from python.core.model_processing.onnx_quantizer.layers.relu import ReluQuantizer
+from python.core.model_processing.onnx_quantizer.layers.sub import SubQuantizer
 
 
 class ONNXOpQuantizer:
@@ -69,6 +74,8 @@ class ONNXOpQuantizer:
 
         # Register handlers
         self.register("Add", AddQuantizer(self.new_initializers))
+        self.register("Sub", SubQuantizer(self.new_initializers))
+        self.register("Mul", MulQuantizer(self.new_initializers))
         self.register("Conv", ConvQuantizer(self.new_initializers))
         self.register("Relu", ReluQuantizer())
         self.register("Reshape", PassthroughQuantizer())
@@ -76,6 +83,7 @@ class ONNXOpQuantizer:
         self.register("Constant", ConstantQuantizer())
         self.register("MaxPool", MaxpoolQuantizer())
         self.register("Flatten", PassthroughQuantizer())
+        self.register("BatchNormalization", BatchnormQuantizer(self.new_initializers))
 
     def register(
         self: ONNXOpQuantizer,
@@ -203,3 +211,32 @@ class ONNXOpQuantizer:
             dict[str, onnx.TensorProto]: Map from initializer name to tensors in graph.
         """
         return {init.name: init for init in model.graph.initializer}
+
+    def apply_pre_analysis_transforms(
+        self: ONNXOpQuantizer,
+        model: onnx.ModelProto,
+        scale_exponent: int,
+        scale_base: int,
+    ) -> onnx.ModelProto:
+        """
+        Give each registered handler a chance to rewrite the model before analysis.
+        """
+        graph = model.graph
+        initializer_map = self.get_initializer_map(model)
+
+        # We allow handlers to modify graph in-place.
+        # (Nodes may be replaced, removed, or new nodes added.)
+        for node in list(graph.node):
+            handler = self.handlers.get(node.op_type)
+            if handler and hasattr(handler, "pre_analysis_transform"):
+                handler.pre_analysis_transform(
+                    node,
+                    graph,
+                    initializer_map,
+                    scale_exponent=scale_exponent,
+                    scale_base=scale_base,
+                )
+            # Refresh map if transforms may add initializers
+            initializer_map = self.get_initializer_map(model)
+
+        return model

python/scripts/gen_and_bench.py

@@ -247,12 +247,12 @@ def export_onnx(
 
 
 def write_input_json(json_path: Path, input_shape: tuple[int] = (1, 4, 28, 28)) -> None:
-    """Write a zero-valued input tensor to JSON alongside its [N,C,H,W] shape."""
+    """Write a zero-valued input tensor to JSON without shape information."""
     json_path.parent.mkdir(parents=True, exist_ok=True)
     n, c, h, w = input_shape
     arr = [0.0] * (n * c * h * w)
     with json_path.open("w", encoding="utf-8") as f:
-        json.dump({"input": arr, "shape": [n, c, h, w]}, f)
+        json.dump({"input": arr}, f)
 
 
 def run_bench(

python/tests/onnx_quantizer_tests/layers/batchnorm_config.py

@@ -0,0 +1,190 @@
+import numpy as np
+
+from python.tests.onnx_quantizer_tests import TEST_RNG_SEED
+from python.tests.onnx_quantizer_tests.layers.base import (
+    BaseLayerConfigProvider,
+    LayerTestConfig,
+    LayerTestSpec,
+    e2e_test,
+    valid_test,
+)
+
+
+class BatchNormConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for BatchNorm (ONNX BatchNormalization op)"""
+
+    @property
+    def layer_name(self) -> str:
+        return "BatchNormalization"
+
+    def get_config(self) -> LayerTestConfig:
+        rng = np.random.default_rng(TEST_RNG_SEED)
+
+        # default shapes: N x C x H x W
+        default_input_shape = [1, 3, 4, 4]
+        c = default_input_shape[1]
+
+        # typical required initializers (scale, bias, mean, var) are length C
+        return LayerTestConfig(
+            op_type="BatchNormalization",
+            valid_inputs=["X", "scale", "B", "input_mean", "input_var"],
+            valid_attributes={
+                "epsilon": 1e-5,
+                "momentum": 0.9,
+                "training_mode": 0,
+            },
+            required_initializers={
+                # Defaults are stored as numpy arrays with shape (C,)
+                "scale": rng.normal(1.0, 0.5, c).astype(np.float32),
+                "B": rng.normal(0.0, 0.5, c).astype(np.float32),
+                "input_mean": rng.normal(0.0, 1.0, c).astype(np.float32),
+                "input_var": np.abs(rng.normal(1.0, 0.5, c)).astype(np.float32),
+            },
+            input_shapes={"X": default_input_shape},
+            output_shapes={"batchnormalization_output": default_input_shape},
+        )
+
+    def get_test_specs(self) -> list[LayerTestSpec]:
+        rng = np.random.default_rng(TEST_RNG_SEED)
+        c = 3
+
+        return [
+            # Basic valid tests
+            valid_test("basic_inference")
+            .description("Basic BatchNormalization inference: standard shapes")
+            .tags("basic", "inference", "batchnorm")
+            .build(),
+            valid_test("different_input_shape")
+            .description("Inference with different spatial dims")
+            .override_input_shapes(X=[2, c, 8, 8])
+            .override_output_shapes(batchnormalization_output=[2, c, 8, 8])
+            .tags("inference", "spatial")
+            .build(),
+            valid_test("epsilon_variation")
+            .description("Inference with larger epsilon for numerical stability")
+            .override_attrs(epsilon=1e-3)
+            .tags("epsilon")
+            .build(),
+            valid_test("momentum_variation")
+            .description(
+                "Inference with non-default momentum (has no effect in inference mode)",
+            )
+            .override_attrs(momentum=0.5)
+            .tags("momentum")
+            .build(),
+            valid_test("zero_mean_input")
+            .description("Input with zero mean")
+            .override_initializer("input_mean", np.zeros((c,), dtype=np.float32))
+            .tags("edge", "zero_mean")
+            .build(),
+            # Scalar / broadcast style tests
+            valid_test("per_channel_zero_variance")
+            .description(
+                "Edge case: very small variance values (clamped by epsilon), inference",
+            )
+            .override_initializer("input_var", np.full((c,), 1e-8, dtype=np.float32))
+            .override_attrs(epsilon=1e-5)
+            .tags("edge", "small_variance")
+            .build(),
+            # E2E tests that set explicit initializer values
+            e2e_test("e2e_inference")
+            .description("E2E inference test with explicit initializers")
+            .override_input_shapes(X=[1, c, 2, 2])
+            .override_output_shapes(batchnormalization_output=[1, c, 2, 2])
+            .override_initializer("scale", rng.normal(1.0, 0.1, c).astype(np.float32))
+            .override_initializer("B", rng.normal(0.0, 0.1, c).astype(np.float32))
+            .override_initializer(
+                "input_mean",
+                rng.normal(0.0, 0.1, c).astype(np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.abs(rng.normal(0.5, 0.2, c)).astype(np.float32),
+            )
+            .tags("e2e", "inference")
+            .build(),
+            e2e_test("e2e_inference_small_2x2")
+            .description("E2E inference with small 2x2 spatial input")
+            .override_input_shapes(X=[1, 3, 2, 2])
+            .override_output_shapes(batchnormalization_output=[1, 3, 2, 2])
+            .override_initializer("scale", np.array([1.0, 0.9, 1.1], dtype=np.float32))
+            .override_initializer("B", np.array([0.0, 0.1, -0.1], dtype=np.float32))
+            .override_initializer(
+                "input_mean",
+                np.array([0.5, -0.5, 0.0], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.array([0.25, 0.5, 0.1], dtype=np.float32),
+            )
+            .tags("e2e", "small", "2x2")
+            .build(),
+            e2e_test("e2e_inference_wide_input")
+            .description("E2E inference with wider input shape (C=4, H=2, W=8)")
+            .override_input_shapes(X=[2, 4, 2, 8])
+            .override_output_shapes(batchnormalization_output=[2, 4, 2, 8])
+            .override_initializer(
+                "scale",
+                np.array([1.0, 0.8, 1.2, 0.9], dtype=np.float32),
+            )
+            .override_initializer(
+                "B",
+                np.array([0.0, 0.1, -0.1, 0.05], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_mean",
+                np.array([0.0, 0.5, -0.5, 0.2], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.array([1.0, 0.5, 0.25, 0.1], dtype=np.float32),
+            )
+            .tags("e2e", "wide", "C4")
+            .build(),
+            e2e_test("e2e_inference_batch2_channels3")
+            .description("E2E inference with batch size 2 and 3 channels")
+            .override_input_shapes(X=[2, 3, 4, 4])
+            .override_output_shapes(batchnormalization_output=[2, 3, 4, 4])
+            .override_initializer("scale", np.array([0.5, 1.0, 1.5], dtype=np.float32))
+            .override_initializer("B", np.array([0.0, 0.0, 0.0], dtype=np.float32))
+            .override_initializer(
+                "input_mean",
+                np.array([-0.5, 0.0, 0.5], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.array([0.2, 0.5, 0.8], dtype=np.float32),
+            )
+            .tags("e2e", "batch2", "C3")
+            .build(),
+            e2e_test("e2e_inference_high_epsilon")
+            .description("E2E inference with high epsilon for numerical stability")
+            .override_input_shapes(X=[1, 2, 4, 4])
+            .override_output_shapes(batchnormalization_output=[1, 2, 4, 4])
+            .override_initializer("scale", np.array([1.0, 1.0], dtype=np.float32))
+            .override_initializer("B", np.array([0.1, -0.1], dtype=np.float32))
+            .override_initializer("input_mean", np.array([0.0, 0.5], dtype=np.float32))
+            .override_initializer(
+                "input_var",
+                np.array([0.0, 0.0], dtype=np.float32),
+            )  # tiny variance
+            .override_attrs(epsilon=1e-2)
+            .tags("e2e", "high_epsilon", "numerical_stability")
+            .build(),
+            e2e_test("e2e_inference_non_square")
+            .description("E2E inference with non-square spatial dimensions")
+            .override_input_shapes(X=[1, 3, 2, 5])
+            .override_output_shapes(batchnormalization_output=[1, 3, 2, 5])
+            .override_initializer("scale", np.array([1.0, 0.9, 1.1], dtype=np.float32))
+            .override_initializer("B", np.array([0.0, 0.1, -0.1], dtype=np.float32))
+            .override_initializer(
+                "input_mean",
+                np.array([0.1, -0.1, 0.0], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.array([0.5, 0.25, 0.75], dtype=np.float32),
+            )
+            .tags("e2e", "non_square", "C3")
+            .build(),
+        ]

python/tests/onnx_quantizer_tests/layers/mul_config.py

@@ -0,0 +1,102 @@
+import numpy as np
+
+from python.tests.onnx_quantizer_tests import TEST_RNG_SEED
+from python.tests.onnx_quantizer_tests.layers.base import (
+    BaseLayerConfigProvider,
+    LayerTestConfig,
+    LayerTestSpec,
+    e2e_test,
+    edge_case_test,
+    valid_test,
+)
+
+
+class MulConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for Mul layer"""
+
+    @property
+    def layer_name(self) -> str:
+        return "Mul"
+
+    def get_config(self) -> LayerTestConfig:
+        return LayerTestConfig(
+            op_type="Mul",
+            valid_inputs=["A", "B"],
+            valid_attributes={},  # Mul has no layer-specific attributes
+            required_initializers={},
+            input_shapes={
+                "A": [1, 3, 4, 4],
+                "B": [1, 3, 4, 4],
+            },
+            output_shapes={
+                "mul_output": [1, 3, 4, 4],
+            },
+        )
+
+    def get_test_specs(self) -> list[LayerTestSpec]:
+        rng = np.random.default_rng(TEST_RNG_SEED)
+        return [
+            # --- VALID TESTS ---
+            valid_test("basic")
+            .description("Basic elementwise Mul of two same-shaped tensors")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 4, 4])
+            .tags("basic", "elementwise", "Mul")
+            .build(),
+            valid_test("broadcast_mul")
+            .description("mul with Numpy-style broadcasting along spatial dimensions")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 1, 1])
+            .tags("broadcast", "elementwise", "mul", "onnx14")
+            .build(),
+            valid_test("initializer_mul")
+            .description(
+                "mul where second input (B) is a tensor initializer instead of input",
+            )
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", rng.normal(0, 1, (1, 3, 4, 4)))
+            .tags("initializer", "elementwise", "mul", "onnxruntime")
+            .build(),
+            valid_test("scalar_mul")
+            .description("mul scalar (initializer) to tensor")
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", np.array([2.0], dtype=np.float32))
+            .tags("scalar", "elementwise", "mul")
+            .build(),
+            # # --- E2E TESTS ---
+            e2e_test("e2e_mul")
+            .description("End-to-end mul test with random inputs")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 4, 4])
+            .override_output_shapes(mul_output=[1, 3, 4, 4])
+            .tags("e2e", "mul", "2d")
+            .build(),
+            e2e_test("e2e_initializer_mul")
+            .description(
+                "mul where second input (B) is a tensor initializer instead of input",
+            )
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", rng.normal(0, 1, (1, 3, 4, 4)))
+            .tags("initializer", "elementwise", "mul", "onnxruntime")
+            .build(),
+            e2e_test("e2e_broadcast_mul")
+            .description("mul with Numpy-style broadcasting along spatial dimensions")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 1, 1])
+            .tags("broadcast", "elementwise", "mul", "onnx14")
+            .build(),
+            e2e_test("e2e_scalar_mul")
+            .description("mul scalar (initializer) to tensor")
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", np.array([2.0], dtype=np.float32))
+            .tags("scalar", "elementwise", "mul")
+            .build(),
+            # # --- EDGE CASES ---
+            edge_case_test("empty_tensor")
+            .description("mul with empty tensor input (zero elements)")
+            .override_input_shapes(A=[0], B=[0])
+            .tags("edge", "empty", "mul")
+            .build(),
+            edge_case_test("large_tensor")
+            .description("Large tensor mul performance/stress test")
+            .override_input_shapes(A=[1, 64, 256, 256], B=[1, 64, 256, 256])
+            .tags("large", "performance", "mul")
+            .skip("Performance test, skipped by default")
+            .build(),
+        ]

python/tests/onnx_quantizer_tests/layers/sub_config.py

@@ -0,0 +1,102 @@
+import numpy as np
+
+from python.tests.onnx_quantizer_tests import TEST_RNG_SEED
+from python.tests.onnx_quantizer_tests.layers.base import (
+    BaseLayerConfigProvider,
+    LayerTestConfig,
+    LayerTestSpec,
+    e2e_test,
+    edge_case_test,
+    valid_test,
+)
+
+
+class SubConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for Sub layer"""
+
+    @property
+    def layer_name(self) -> str:
+        return "Sub"
+
+    def get_config(self) -> LayerTestConfig:
+        return LayerTestConfig(
+            op_type="Sub",
+            valid_inputs=["A", "B"],
+            valid_attributes={},  # Sub has no layer-specific attributes
+            required_initializers={},
+            input_shapes={
+                "A": [1, 3, 4, 4],
+                "B": [1, 3, 4, 4],
+            },
+            output_shapes={
+                "sub_output": [1, 3, 4, 4],
+            },
+        )
+
+    def get_test_specs(self) -> list[LayerTestSpec]:
+        rng = np.random.default_rng(TEST_RNG_SEED)
+        return [
+            # --- VALID TESTS ---
+            valid_test("basic")
+            .description("Basic elementwise Sub of two same-shaped tensors")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 4, 4])
+            .tags("basic", "elementwise", "Sub")
+            .build(),
+            valid_test("broadcast_Sub")
+            .description("Sub with Numpy-style broadcasting along spatial dimensions")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 1, 1])
+            .tags("broadcast", "elementwise", "Sub", "onnx14")
+            .build(),
+            valid_test("initializer_Sub")
+            .description(
+                "Sub where second input (B) is a tensor initializer instead of input",
+            )
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", rng.normal(0, 1, (1, 3, 4, 4)))
+            .tags("initializer", "elementwise", "Sub", "onnxruntime")
+            .build(),
+            valid_test("scalar_Sub")
+            .description("Sub scalar (initializer) to tensor")
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", np.array([2.0], dtype=np.float32))
+            .tags("scalar", "elementwise", "Sub")
+            .build(),
+            # --- E2E TESTS ---
+            e2e_test("e2e_Sub")
+            .description("End-to-end Sub test with random inputs")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 4, 4])
+            .override_output_shapes(sub_output=[1, 3, 4, 4])
+            .tags("e2e", "Sub", "2d")
+            .build(),
+            e2e_test("e2e_initializer_Sub")
+            .description(
+                "Sub where second input (B) is a tensor initializer instead of input",
+            )
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", rng.normal(0, 1, (1, 3, 4, 4)))
+            .tags("initializer", "elementwise", "Sub", "onnxruntime")
+            .build(),
+            e2e_test("e2e_broadcast_Sub")
+            .description("Sub with Numpy-style broadcasting along spatial dimensions")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 1, 1])
+            .tags("broadcast", "elementwise", "Sub", "onnx14")
+            .build(),
+            e2e_test("e2e_scalar_Sub")
+            .description("Sub scalar (initializer) to tensor")
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", np.array([2.0], dtype=np.float32))
+            .tags("scalar", "elementwise", "Sub")
+            .build(),
+            # # --- EDGE CASES ---
+            edge_case_test("empty_tensor")
+            .description("Sub with empty tensor input (zero elements)")
+            .override_input_shapes(A=[0], B=[0])
+            .tags("edge", "empty", "Sub")
+            .build(),
+            edge_case_test("large_tensor")
+            .description("Large tensor Sub performance/stress test")
+            .override_input_shapes(A=[1, 64, 256, 256], B=[1, 64, 256, 256])
+            .tags("large", "performance", "Sub")
+            .skip("Performance test, skipped by default")
+            .build(),
+        ]

python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py

@@ -139,6 +139,8 @@ class TestQuantize(BaseQuantizerTest):
             node: NodeProto,
             result_node: NodeProto,
         ) -> bool:
+            if node.op_type == "BatchNormalization":
+                pytest.skip(f"{node.op_type} alters the node structure by design")
             if node.op_type in result_node.op_type:
                 # Assert there are no less attributes in the new node
                 assert len(node.attribute) <= len(result_node.attribute)