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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: JSTprove
-Version: 1.0.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**.
+* ✅ 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.0.0.dist-info → jstprove-1.2.0.dist-info}/RECORD

@@ -1,47 +1,53 @@
-jstprove-1.0.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-0-0,sha256=Lh1hGumXMjFTJEUD9SeYY2RIMl635L35Y7W7GpR0SQM,3027264
+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=OevM274yeaMQIroAv-eTjUcu_tY1Kr9BqEKdoEIpKjM,7597
+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=GNHJs7FBGJZdvlqGxsLix2TYc41hs-jsK23gSAFWyyU,35951
+python/core/circuits/base.py,sha256=tvCHwk_V2ftEocQkmoK5Nf98Iy0F0Ce1FWp03HhNHfA,41274
 python/core/circuits/errors.py,sha256=KzIXyi2ssVvBmXV0Rgn0dBfsTgweKHjeSvP2byRmqGc,5964
 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=ks78fNZfgHU0RDnZZ2T3cdrlDxLnk7jqeR9o8OSqfFw,42862
+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=dlnFIRX3lZGtpqmwSqmhRaCJpogefdtXEAUlv4NuJ0U,2316
+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=iEOnGn9Hoipta3UepKIz5oO1veBdcctcLRRT-_xMPqE,1202
+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=JpTZZKlMRLavm7vM8VNC1XQY4up7rY5Je7XIezjI1fc,7468
+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/base.py,sha256=wY2_65h0bPJHKErGlC4zamSKhhBKq3P7BEYIgwPKAV8,13259
-python/core/model_processing/onnx_quantizer/layers/constant.py,sha256=iWfurTtJk-KaQW6OI4oIDvA8H3clfVwTmf5Xy-QrM0U,3699
-python/core/model_processing/onnx_quantizer/layers/conv.py,sha256=ay9pkctbPfHmra9qZLhz5YCcEGuI67lnXZbJ5rNTgH8,6043
-python/core/model_processing/onnx_quantizer/layers/gemm.py,sha256=VRhN2goyhKFUmGMFX14p3vHvxNAvA1GLKyzjCRUnBoI,5554
-python/core/model_processing/onnx_quantizer/layers/maxpool.py,sha256=ie9rAwdQAh-QoymkrkEeFcRNLqgVk2MdmuJs4iRTc8k,4650
-python/core/model_processing/onnx_quantizer/layers/relu.py,sha256=h6KLK4rboTMc3ZA6GpsLTLkD74pumshaA8igLQeb_pw,2357
+python/core/model_processing/onnx_quantizer/layers/add.py,sha256=AGxzqMa0jABIEKOIgPqEAA7EpZtynQtnD9nxI2NHc0s,1409
+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
 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/general_layer_functions.py,sha256=l_aIqMbSUKxWdUTI_sVLqTdNZj_NnpzXXejSHidhoag,9706
+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=nhx1N2bV6YtqEtRuCnA6GB8TSZf2kdwzuGlUgYQc9F8,4789
+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/witness_utils.py,sha256=ukvbF6EaHMPzRQVZad9wQ9gISRwBGQ1hEAHzc5TpGuw,9488
 python/frontend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -61,21 +67,50 @@ 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/other_e2e_test.py,sha256=o2YfYNwdWRbalB4pE3qZTyTs5nHAAWN33vVHtl_OHv0,6034
+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=t4h4Bkz3PJpm2fARGw6TJS4S9cE0uTa8FFkkAUaBGEI,30685
-python/tests/circuit_parent_classes/test_onnx_converter.py,sha256=naNDsxpcxH2ClsUgv5_TS_ZOpLo8XSGse_x7gfFLTUY,6266
+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/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/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/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=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.0.0.dist-info/METADATA,sha256=nM_w5Qsh8TwJJ6PDGZcBzND029Fj4uWwmilqQXVx_mw,14047
-jstprove-1.0.0.dist-info/WHEEL,sha256=jc2C2uw104ioj1TL9cE0YO67_kdAwX4W8JgYPomxr5M,105
-jstprove-1.0.0.dist-info/entry_points.txt,sha256=nGcTSO-4q08gPl1IoWdrPaiY7IbO7XvmXKkd34dYHc8,49
-jstprove-1.0.0.dist-info/top_level.txt,sha256=J-z0poNcsv31IHB413--iOY8LoHBKiTHeybHX3abokI,7
-jstprove-1.0.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/circuit_models/generic_onnx.py

@@ -127,11 +127,16 @@ class GenericModelONNX(ONNXConverter, ZKModelBase):
             )
         return models_onnx_path
 
-    def adjust_inputs(self: GenericModelONNX, input_file: str) -> str:
+    def adjust_inputs(
+        self: GenericModelONNX,
+        inputs: dict[str, np.ndarray],
+        input_file: str,
+    ) -> str:
         """Preprocess and flatten model inputs for the circuit.
 
         Args:
-            input_file (str): Input data file or array compatible with the model.
+            inputs (str): inputs, read from json file
+            input_file (str): path to input_file
 
         Returns:
             str: Adjusted input file after reshaping and scaling.
@@ -140,7 +145,7 @@ class GenericModelONNX(ONNXConverter, ZKModelBase):
             input_shape = self.input_shape.copy()
             shape = self.adjust_shape(input_shape)
             self.input_shape = [math.prod(shape)]
-            x = super().adjust_inputs(input_file)
+            x = super().adjust_inputs(inputs, input_file)
             self.input_shape = input_shape.copy()
         except Exception as e:
             msg = f"Failed to adjust inputs for GenericModelONNX: {e}"
@@ -169,7 +174,9 @@ class GenericModelONNX(ONNXConverter, ZKModelBase):
                 operation="get_outputs",
             ) from e
         else:
-            return torch.as_tensor(np.array(raw_outputs)).flatten()
+            flat_outputs = [o.flatten() for o in raw_outputs]
+            combined = np.concatenate(flat_outputs, axis=0)
+            return torch.as_tensor(combined)
 
     def format_inputs(
         self: GenericModelONNX,
@@ -186,15 +193,42 @@ class GenericModelONNX(ONNXConverter, ZKModelBase):
         Returns:
             Dict[str, List[int]]: Dictionary mapping `input` to scaled integer values.
         """
+
+        def _raise_type_error(inputs: np.ndarray | list[int] | torch.Tensor) -> None:
+            msg = (
+                "Expected np.ndarray, torch.Tensor, "
+                f"list, or dict for inputs, got {type(inputs)}"
+            )
+
+            raise TypeError(msg)
+
         try:
-            x = {"input": inputs}
+            if isinstance(inputs, (np.ndarray, torch.Tensor, list)):
+                inputs = {"input": inputs}
+            elif not isinstance(inputs, dict):
+                _raise_type_error(inputs)
             scaling = BaseOpQuantizer.get_scaling(
                 scale_base=self.scale_base,
                 scale_exponent=self.scale_exponent,
             )
-            for key in x:  # noqa: PLC0206
-                x[key] = torch.as_tensor(x[key]).flatten().tolist()
-                x[key] = (torch.as_tensor(x[key]) * scaling).long().tolist()
+
+            input_shapes: dict[str, tuple[int, ...]] = {}
+            flattened_tensors: list[torch.Tensor] = []
+
+            # Flatten, scale, and collect each input
+            for name, value in inputs.items():
+                tensor = torch.as_tensor(value)
+                input_shapes[name] = tuple(tensor.shape)
+
+                scaled = (tensor * scaling).long().flatten()
+                flattened_tensors.append(scaled)
+
+            # Concatenate all inputs into one long tensor
+            concatenated = torch.cat(flattened_tensors, dim=0)
+            flattened_list = concatenated.tolist()
+
+            # Wrap it into a dict under "input" key to read into rust
+            formatted_inputs = {"input": flattened_list}
         except Exception as e:
             msg = f"Failed to format inputs for GenericModelONNX: {e}"
             raise CircuitProcessingError(
@@ -203,7 +237,7 @@ class GenericModelONNX(ONNXConverter, ZKModelBase):
                 data_type=type(inputs).__name__,
             ) from e
         else:
-            return x
+            return formatted_inputs
 
     def get_weights(
         self: GenericModelONNX,