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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: JSTprove
-Version: 1.1.0
+Version: 1.3.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.3.0.dist-info}/RECORD

@@ -1,49 +1,57 @@
-jstprove-1.1.0.dist-info/licenses/LICENSE,sha256=UXQRcYRUH-PfN27n3P-FMaZFY6jr9jFPKcwT7CWbljw,1160
+jstprove-1.3.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-3-0,sha256=qbHC9SV_NuNv-vZs5MwXV0NxXkqTTZlmLH59WYCVrC8,3221088
 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
 python/core/circuits/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-python/core/circuits/base.py,sha256=tvCHwk_V2ftEocQkmoK5Nf98Iy0F0Ce1FWp03HhNHfA,41274
-python/core/circuits/errors.py,sha256=KzIXyi2ssVvBmXV0Rgn0dBfsTgweKHjeSvP2byRmqGc,5964
+python/core/circuits/base.py,sha256=_XSs2LFyBMZEkBSvRp53zc-XaGOopNxYg2xHQz9sqt0,41991
+python/core/circuits/errors.py,sha256=JDNa23wMwNQDTFY0IpDpHDMZ9gOdjDdQmB4GBhL_DCg,5913
 python/core/circuits/zk_model_base.py,sha256=5ggOaJjs2_MJvn-PO1cPN3i7U-XR4L-0zJGYuLVKOLc,820
 python/core/model_processing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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/onnx_custom_ops/__init__.py,sha256=ofecV9pzpDJJl_r6inRw9JOKxtfK2rzzxWahAq9BKXE,475
+python/core/model_processing/converters/base.py,sha256=o6bNwmqD9sOM9taqMb0ed6804RugQiU3va0rY_EA5SE,4265
+python/core/model_processing/converters/onnx_converter.py,sha256=-eXdF6tfluFRxGgnQtJQ8R2309aYX-8z8HzMxk_Qv8I,44340
+python/core/model_processing/onnx_custom_ops/__init__.py,sha256=ZKUC4ToRxgEEMHcTyERATVEN0KSDs-9cM1T-tTw3I1g,525
+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/exceptions.py,sha256=vzxBRbpvk4ZZbgacDISnqmQQKj7Ls46V08ilHnhaJy0,5645
+python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py,sha256=5I67frJn4j2T1LTvODHixQK4VaqazJFJ0T1BCvqLPgg,9655
 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=zUUAZpXCtbxbhHzuYczTNZPe-xWr6TxpmoAIDe4kCo4,21176
+python/core/model_processing/onnx_quantizer/layers/batchnorm.py,sha256=KSBDPHd52f5Qyf-cnIDFPmfzssaJgMPiTmpIWEdM41U,7718
+python/core/model_processing/onnx_quantizer/layers/clip.py,sha256=HrhiLtqC3cIAvU0wRCqp8_8ZSFH8a3F1Jf_qkXlY44s,3043
 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/max.py,sha256=3gUxrdXwcVAtgR-_j4xQ0085Wj0oEBLT897TImxF2d4,1343
 python/core/model_processing/onnx_quantizer/layers/maxpool.py,sha256=PJ8hZPPBpfWV_RZdySl50-BU8TATjcg8Tg_mrAVS1Ic,4916
+python/core/model_processing/onnx_quantizer/layers/min.py,sha256=cQbXzGOApR6HUJZMARXy87W8IbUC562jnAQm8J8ynQI,1709
+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/model_templates/circuit_template.py,sha256=OAqMRshi9OiJYoqpjkg5tUfNf18MfZmhsxxD6SANm_4,2106
 python/core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/core/utils/benchmarking_helpers.py,sha256=0nT38SCrjP_BlvJODsc9twF9ZmIFg_1sAvSyeNfv4mQ,5235
 python/core/utils/constants.py,sha256=Qu5_6OUe1XIsL-IY5_4923eN7x1-SPv6ohQonztAobA,102
-python/core/utils/errors.py,sha256=vTlluhbSqmyI5e1JNLEZ1mQ-dG_Wbxe4p5l4aa59zAY,3739
+python/core/utils/errors.py,sha256=Uf57cRKpot_u5Yr8HRmjLmInkdd_x0x5YpTGBncZgl4,3722
 python/core/utils/general_layer_functions.py,sha256=tg2WWhmR-4TlKn8OeCu1qNbLf8qdKVP3jl9mhZn_sTg,9781
 python/core/utils/helper_functions.py,sha256=3JwJa4wHoUBteukDw4bAetqMsQLeJ0_sJ0qIdKy7GCY,37097
 python/core/utils/model_registry.py,sha256=aZg_9LEqsBXK84oxQ8A3NGZl-9aGnLgfR-kgxkOwV50,4895
-python/core/utils/scratch_tests.py,sha256=UYXsWIBh_27OxnyfH9CuxeNFT-OWCK0YpJ-j-8f0QHc,2332
+python/core/utils/scratch_tests.py,sha256=o2VDTk8QBKA3UHHE-h7Ghtoge6kGG7G-8qwvesuTFFc,2281
 python/core/utils/witness_utils.py,sha256=ukvbF6EaHMPzRQVZad9wQ9gISRwBGQ1hEAHzc5TpGuw,9488
 python/frontend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/frontend/cli.py,sha256=lkvhzQC6bv0AgWUypg_cH-JT574r89qgTIsgHDT9GRg,3106
@@ -62,47 +70,53 @@ 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
-python/tests/circuit_e2e_tests/circuit_model_developer_test.py,sha256=Ic9hprCn1Rs-XAF-SUmBNEDn65yaCxUK9z5875KPg5o,39416
-python/tests/circuit_e2e_tests/helper_fns_for_tests.py,sha256=4toXD0pJMYVZFL1O9JZAJF-iqbi9k1eyuk_goUnchRo,5190
+python/tests/circuit_e2e_tests/circuit_model_developer_test.py,sha256=8hl8SKw7obXplo0jsiKoKIZLxlu1_HhXvGDeSBDBars,39456
+python/tests/circuit_e2e_tests/helper_fns_for_tests.py,sha256=uEThqTsRdNJivHwAv-aJIUtSPlmVHdhMZqZSH1OqhDE,5177
 python/tests/circuit_e2e_tests/other_e2e_test.py,sha256=amWRa1tIBHdQpd9-XS7vBXG0tkdV_9K9fH-FT5LFh7E,11301
 python/tests/circuit_parent_classes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 python/tests/circuit_parent_classes/test_circuit.py,sha256=5vgcZHD2wY_pIRFNAhEZuBJD4uw2QyTck75Z9CJaACE,45968
 python/tests/circuit_parent_classes/test_onnx_converter.py,sha256=sJ0o8sducNUtmYKmsqfx7WEsIEd6oNbnWk71rXS_nIU,6575
-python/tests/circuit_parent_classes/test_ort_custom_layers.py,sha256=FEEY8nbuMC2xb6WrBsik7TeDde6SlMwwG9PKSqmCymo,3980
+python/tests/circuit_parent_classes/test_ort_custom_layers.py,sha256=PBKjt6Mu3lRco4ijD2BLwAHPWRFic-OUwWPVsvBoEpU,3042
 python/tests/onnx_quantizer_tests/__init__.py,sha256=IZPGWHgjoay3gM1p2WJNh5cnZ79EP2VP-bcKy8AfJjY,18
 python/tests/onnx_quantizer_tests/test_base_layer.py,sha256=Ro7k-eUbGCyfIZ-OVNjLlCIz3mb02uHFWboFuWOdXKs,6526
 python/tests/onnx_quantizer_tests/test_exceptions.py,sha256=pwhARalEXx7REkcnIVZPi-4J1wgzgZN4xG-wLsx4rTs,3473
 python/tests/onnx_quantizer_tests/test_onnx_op_quantizer.py,sha256=m6mNe1KDRFIE2P0YURTIAim9-Di0BoPPAaaOOlorDIk,7367
-python/tests/onnx_quantizer_tests/test_registered_quantizers.py,sha256=M8N3KxApfIpZIu2Swh_z8eSy3DDqB3XxebN685hHHlw,4052
+python/tests/onnx_quantizer_tests/test_registered_quantizers.py,sha256=lw_jYSbQ9ZM9P-jSt_LFSuve9vQ22cLtSui0W3zGqpo,4209
 python/tests/onnx_quantizer_tests/testing_helper_functions.py,sha256=N0fQv2pYzUCVZ7wkcR8gEKs5zTXT1hWrK-HKSTQYvYU,534
 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/base.py,sha256=3nqmU2PgOdK_mPkz-YHg3idgr-PXYbu5kCIY-Uic5yo,9317
+python/tests/onnx_quantizer_tests/layers/batchnorm_config.py,sha256=P-sZuHAdEfNczcgTeLjqJnEbpqN3dKTsbqvY4-SBqiQ,8231
+python/tests/onnx_quantizer_tests/layers/clip_config.py,sha256=-OuhnUgz6xY4iW1jUR7W-J__Ie9lXI9vplmzp8qXqRc,4973
 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/max_config.py,sha256=vzR8-2wbPGcH0GMmAJ_sXSEdMtZOjVNGufU__N3Jfyw,3906
 python/tests/onnx_quantizer_tests/layers/maxpool_config.py,sha256=XfTPk_ZQXEzaCjHHymSLVv2HS-PKH1rS9IuyyoEtM78,3176
+python/tests/onnx_quantizer_tests/layers/min_config.py,sha256=izKtCaMXoQHiAfmcGlJRQdKMQz3Su8n0p2mEn0y56Do,3774
+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_integration.py,sha256=xNt2STeXB33NcpteDThwGTSW1Hm15POf8a4aPBSVrvI,7254
+python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py,sha256=DatmgvibQazP100B4NHDu7u-O2-f90juPKvPOXuPnXo,9491
 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.3.0.dist-info/METADATA,sha256=CqGuzrQWy_MUYtOcy-0i8mB_2eAKUvQ1R3tjEX-3N4o,14100
+jstprove-1.3.0.dist-info/WHEEL,sha256=jc2C2uw104ioj1TL9cE0YO67_kdAwX4W8JgYPomxr5M,105
+jstprove-1.3.0.dist-info/entry_points.txt,sha256=nGcTSO-4q08gPl1IoWdrPaiY7IbO7XvmXKkd34dYHc8,49
+jstprove-1.3.0.dist-info/top_level.txt,sha256=J-z0poNcsv31IHB413--iOY8LoHBKiTHeybHX3abokI,7
+jstprove-1.3.0.dist-info/RECORD,,

python/core/circuits/base.py

@@ -4,13 +4,12 @@ import logging
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-from numpy import asarray, ndarray
+import numpy as np
 
 from python.core.utils.errors import ShapeMismatchError
 from python.core.utils.witness_utils import compare_witness_to_io, load_witness
 
 if TYPE_CHECKING:
-    import numpy as np
     import torch
 
 from python.core.circuits.errors import (
@@ -775,18 +774,18 @@ class Circuit:
     def reshape_inputs_for_inference(
         self: Circuit,
         inputs: dict[str],
-    ) -> ndarray | dict[str, ndarray]:
+    ) -> np.ndarray | dict[str, np.ndarray]:
         """
         Reshape input tensors to match the model's expected input shape.
 
         Parameters
         ----------
-        inputs : dict[str] or ndarray
+        inputs : dict[str] or np.ndarray
             Input tensors or a dictionary of tensors.
 
         Returns
         -------
-        ndarray or dict[str, ndarray]
+        np.ndarray or dict[str, np.ndarray]
             Reshaped input(s) ready for inference.
         """
 
@@ -801,15 +800,33 @@ class Circuit:
         if isinstance(inputs, dict):
             if len(inputs) == 1:
                 only_key = next(iter(inputs))
-                inputs = asarray(inputs[only_key])
+                value = np.asarray(inputs[only_key])
+
+                # If shape is a dict, extract the shape for this key
+                if isinstance(shape, dict):
+                    key_shape = shape.get(only_key, None)
+                    if key_shape is None:
+                        raise CircuitConfigurationError(
+                            missing_attributes=[f"input_shape[{only_key!r}]"],
+                        )
+                    shape = key_shape
+
+                # From here on, treat it as a regular reshape
+                inputs = value
             else:
                 return self._reshape_dict_inputs(inputs, shape)
 
         # --- Regular reshape ---
+        if not isinstance(shape, (list, tuple)):
+            msg = (
+                f"Expected list or tuple shape for reshape, got {type(shape).__name__}"
+            )
+            raise CircuitInputError(msg)
+
         try:
-            return asarray(inputs).reshape(shape)
+            return np.asarray(inputs).reshape(shape)
         except Exception as e:
-            raise ShapeMismatchError(shape, list(asarray(inputs).shape)) from e
+            raise ShapeMismatchError(shape, list(np.asarray(inputs).shape)) from e
 
     def _reshape_dict_inputs(
         self: Circuit,
@@ -824,7 +841,7 @@ class Circuit:
             )
             raise CircuitInputError(msg, parameter="shape", expected="dict")
         for key, value in inputs.items():
-            tensor = asarray(value)
+            tensor = np.asarray(value)
             try:
                 inputs[key] = tensor.reshape(shape[key])
             except Exception as e:
@@ -867,16 +884,16 @@ class Circuit:
             value = inputs[key]
 
             # --- handle unsupported input types BEFORE entering try ---
-            if not isinstance(value, (ndarray, list, tuple)):
+            if not isinstance(value, (np.ndarray, list, tuple)):
                 msg = f"Unsupported input type for key '{key}': {type(value).__name__}"
                 raise CircuitProcessingError(message=msg)
 
             try:
                 # Convert to tensor, flatten, and back to list
-                if isinstance(value, ndarray):
+                if isinstance(value, np.ndarray):
                     flattened = value.flatten().tolist()
                 else:
-                    flattened = asarray(value).flatten().tolist()
+                    flattened = np.asarray(value).flatten().tolist()
             except Exception as e:
                 msg = f"Failed to flatten input '{key}' (type {type(value).__name__})"
                 raise CircuitProcessingError(message=msg) from e

python/core/circuits/errors.py

@@ -1,4 +1,3 @@
-# python/core/utils/exceptions.py
 from __future__ import annotations
 
 from python.core.utils.helper_functions import RunType
@@ -68,7 +67,7 @@ class CircuitInputError(CircuitError):
         actual (any): Actual value encountered (optional).
     """
 
-    def __init__(  # noqa: PLR0913
+    def __init__(
         self: CircuitInputError,
         message: str | None = None,
         parameter: str | None = None,

python/core/model_processing/converters/base.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     import numpy as np
@@ -16,10 +16,10 @@ class ModelType(Enum):
 
 ONNXLayerDict = dict[
     str,
-    Union[int, str, list[str], dict[str, list[int]], Optional[list], Optional[dict]],
+    int | str | list[str] | dict[str, list[int]] | list | None | dict,
 ]
 
-CircuitParamsDict = dict[str, Union[int, dict[str, bool]]]
+CircuitParamsDict = dict[str, int | dict[str, bool]]
 
 
 class ModelConverter(ABC):

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/__init__.py

@@ -1,16 +1,17 @@
 import importlib
 import pkgutil
-import os
+from pathlib import Path
 
 # Get the package name of the current module
 package_name = __name__
 
 # Dynamically import all .py files in this package directory (except __init__.py)
-package_dir = os.path.dirname(__file__)
+package_dir = Path(__file__).parent.as_posix()
 
-__all__ = []
+
+__all__: list[str] = []
 
 for _, module_name, is_pkg in pkgutil.iter_modules([package_dir]):
     if not is_pkg and (module_name != "custom_helpers"):
         importlib.import_module(f"{package_name}.{module_name}")
-        __all__.append(module_name)
+        __all__.append(module_name)  # noqa: PYI056

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/exceptions.py

@@ -31,7 +31,7 @@ class InvalidParamError(QuantizationError):
     quantization the quantization process.
     """
 
-    def __init__(  # noqa: PLR0913
+    def __init__(
         self: QuantizationError,
         node_name: str,
         op_type: str,
@@ -151,7 +151,7 @@ class InvalidConfigError(QuantizationError):
     def __init__(
         self: QuantizationError,
         key: str,
-        value: str | float | bool | None,
+        value: str | float | bool | None,  # noqa: FBT001
         expected: str | None = None,
     ) -> None:
         """Initialize InvalidConfigError with context about the bad config.