chore: update tract (#766 )

fix: patch large batch ex (#763 )
fix: make commitment optional for backwards compat (#762 )
2026-01-13 16:27:59 -05:00 · 2024-04-04 18:07:08 +01:00 · 2024-04-03 02:33:57 +01:00 · 2024-04-03 02:26:50 +01:00 · 2024-04-02 21:37:20 +01:00 · 2024-04-02 17:08:26 +01:00
53 changed files with 5088 additions and 732 deletions
--- a/.github/workflows/pypi-gpu.yml
+++ b/.github/workflows/pypi-gpu.yml
@@ -26,7 +26,7 @@ jobs:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: 3.7
+          python-version: 3.12
          architecture: x64

      - name: Set pyproject.toml version to match github tag
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -25,7 +25,7 @@ jobs:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: 3.7
+          python-version: 3.12
          architecture: x64

      - name: Set Cargo.toml version to match github tag
@@ -70,7 +70,7 @@ jobs:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: 3.7
+          python-version: 3.12
          architecture: ${{ matrix.target }}

      - name: Set Cargo.toml version to match github tag
@@ -115,7 +115,7 @@ jobs:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: 3.7
+          python-version: 3.12
          architecture: x64

      - name: Set Cargo.toml version to match github tag
@@ -139,6 +139,20 @@ jobs:
          target: ${{ matrix.target }}
          manylinux: auto
          args: --release --out dist --features python-bindings
+          before-script-linux: |
+            # If we're running on rhel centos, install needed packages.
+            if command -v yum &> /dev/null; then
+                yum update -y && yum install -y perl-core openssl openssl-devel pkgconfig libatomic
+
+                # If we're running on i686 we need to symlink libatomic
+                # in order to build openssl with -latomic flag.
+                if [[ ! -d "/usr/lib64" ]]; then
+                    ln -s /usr/lib/libatomic.so.1 /usr/lib/libatomic.so
+                fi
+            else
+                # If we're running on debian-based system.
+                apt update -y && apt-get install -y libssl-dev openssl pkg-config
+            fi

      - name: Install built wheel
        if: matrix.target == 'x86_64'
@@ -162,7 +176,7 @@ jobs:
  #   - uses: actions/checkout@v4
  #   - uses: actions/setup-python@v4
  #     with:
-  #       python-version: 3.7
+  #       python-version: 3.12

  #   - name: Install cross-compilation tools for aarch64
  #     if: matrix.target == 'aarch64'
@@ -214,7 +228,7 @@ jobs:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: 3.7
+          python-version: 3.12
          architecture: x64

      - name: Set Cargo.toml version to match github tag
@@ -249,7 +263,7 @@ jobs:
            apk add py3-pip
            pip3 install -U pip
            python3 -m venv .venv
-            source .venv/bin/activate            
+            source .venv/bin/activate
            pip3 install ezkl --no-index --find-links /io/dist/ --force-reinstall
            python3 -c "import ezkl"

@@ -273,7 +287,7 @@ jobs:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: 3.7
+          python-version: 3.12

      - name: Set Cargo.toml version to match github tag
        shell: bash
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -184,7 +184,7 @@ jobs:

  wasm32-tests:
    runs-on: ubuntu-latest
-    # needs: [build,  library-tests, docs]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions-rs/toolchain@v1
@@ -207,7 +207,7 @@ jobs:

  tutorial:
    runs-on: ubuntu-latest
-    needs: [build, library-tests, docs]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions-rs/toolchain@v1
@@ -224,7 +224,7 @@ jobs:

  mock-proving-tests:
    runs-on: non-gpu
-    # needs: [build,  library-tests, docs]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions-rs/toolchain@v1
@@ -281,7 +281,7 @@ jobs:

  prove-and-verify-evm-tests:
    runs-on: non-gpu
-    needs: [build, library-tests]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions-rs/toolchain@v1
@@ -354,7 +354,7 @@ jobs:

  prove-and-verify-tests:
    runs-on: non-gpu
-    needs: [build, library-tests]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions-rs/toolchain@v1
@@ -460,7 +460,7 @@ jobs:

  prove-and-verify-mock-aggr-tests:
    runs-on: self-hosted
-    needs: [build, library-tests]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions-rs/toolchain@v1
@@ -495,7 +495,7 @@ jobs:

  prove-and-verify-aggr-tests:
    runs-on: large-self-hosted
-    needs: [build, library-tests]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions-rs/toolchain@v1
@@ -512,7 +512,7 @@ jobs:

  prove-and-verify-aggr-evm-tests:
    runs-on: large-self-hosted
-    needs: [build, library-tests]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions-rs/toolchain@v1
@@ -557,16 +557,18 @@ jobs:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: "3.7"
+          python-version: "3.12"
      - uses: actions-rs/toolchain@v1
        with:
          toolchain: nightly-2024-02-06
          override: true
          components: rustfmt, clippy
+      - name: Install cmake
+        run: sudo apt-get install -y cmake
      - name: Install solc
        run: (hash svm 2>/dev/null || cargo install svm-rs) && svm install 0.8.20 && solc --version
      - name: Setup Virtual Env and Install python dependencies
-        run: python -m venv .env; source .env/bin/activate; pip install -r requirements.txt;
+        run: python -m venv .env --clear; source .env/bin/activate; pip install -r requirements.txt;
      - name: Install Anvil
        run: cargo install --git https://github.com/foundry-rs/foundry --rev c2233ec9fe61e0920c61c6d779bc707252852037 --profile local --locked anvil --force
      - name: Build python ezkl
@@ -576,12 +578,12 @@ jobs:

  accuracy-measurement-tests:
    runs-on: ubuntu-latest-32-cores
-    # needs: [build, library-tests, docs]
+    needs: [build, library-tests, docs, python-tests, python-integration-tests]
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: "3.7"
+          python-version: "3.12"
      - uses: actions-rs/toolchain@v1
        with:
          toolchain: nightly-2024-02-06
@@ -592,7 +594,7 @@ jobs:
          crate: cargo-nextest
          locked: true
      - name: Setup Virtual Env and Install python dependencies
-        run: python -m venv .env; source .env/bin/activate; pip install -r requirements.txt;
+        run: python -m venv .env --clear; source .env/bin/activate; pip install -r requirements.txt;
      - name: Build python ezkl
        run: source .env/bin/activate; unset CONDA_PREFIX; maturin develop --features python-bindings --release
      - name: Div rebase
@@ -612,7 +614,7 @@ jobs:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
-          python-version: "3.10"
+          python-version: "3.11"
      - uses: actions-rs/toolchain@v1
        with:
          toolchain: nightly-2024-02-06
@@ -626,10 +628,14 @@ jobs:
        run: (hash svm 2>/dev/null || cargo install svm-rs) && svm install 0.8.20 && solc --version
      - name: Install Anvil
        run: cargo install --git https://github.com/foundry-rs/foundry --rev c2233ec9fe61e0920c61c6d779bc707252852037 --profile local --locked anvil --force
+      - name: Install pip
+        run: python -m ensurepip --upgrade
      - name: Setup Virtual Env and Install python dependencies
-        run: python -m venv .env; source .env/bin/activate; pip install -r requirements.txt;
+        run: python -m venv .env --clear; source .env/bin/activate; pip install -r requirements.txt; python -m ensurepip --upgrade
      - name: Build python ezkl
        run: source .env/bin/activate; unset CONDA_PREFIX; maturin develop --features python-bindings --release
+      - name: Tictactoe tutorials
+        run: source .env/bin/activate; cargo nextest run py_tests::tests::tictactoe_ --test-threads 1
      # - name: authenticate-kaggle-cli
      #   shell: bash
      #   env:
@@ -645,7 +651,5 @@ jobs:
        run: source .env/bin/activate; cargo nextest run py_tests::tests::voice_
      - name: NBEATS tutorial
        run: source .env/bin/activate; cargo nextest run py_tests::tests::nbeats_
-      - name: Tictactoe tutorials
-        run: source .env/bin/activate; cargo nextest run py_tests::tests::tictactoe_
      # - name: Postgres tutorials
      #   run: source .env/bin/activate; cargo nextest run py_tests::tests::postgres_ --test-threads 1
--- a/.gitignore
+++ b/.gitignore
@@ -48,4 +48,5 @@ node_modules
 /dist
 timingData.json
 !tests/wasm/pk.key
-!tests/wasm/vk.key
+!tests/wasm/vk.key
+!tests/wasm/vk_aggr.key
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -112,7 +112,7 @@ checksum = "1a047897373be4bbb0224c1afdabca92648dc57a9c9ef6e7b0be3aff7a859c83"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -211,9 +211,9 @@ checksum = "d301b3b94cb4b2f23d7917810addbbaff90738e0ca2be692bd027e70d7e0330c"

 [[package]]
 name = "arc-swap"
-version = "1.7.0"
+version = "1.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7b3d0060af21e8d11a926981cc00c6c1541aa91dd64b9f881985c3da1094425f"
+checksum = "69f7f8c3906b62b754cd5326047894316021dcfe5a194c8ea52bdd94934a3457"

 [[package]]
 name = "ark-bls12-377"
@@ -463,7 +463,7 @@ dependencies = [
 "proc-macro2",
 "quote",
 "serde",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -483,13 +483,13 @@ dependencies = [

 [[package]]
 name = "async-trait"
-version = "0.1.78"
+version = "0.1.79"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "461abc97219de0eaaf81fe3ef974a540158f3d079c2ab200f891f1a2ef201e85"
+checksum = "a507401cad91ec6a857ed5513a2073c82a9b9048762b885bb98655b306964681"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -522,20 +522,20 @@ checksum = "3c87f3f15e7794432337fc718554eaa4dc8f04c9677a950ffe366f20a162ae42"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
 name = "autocfg"
-version = "1.1.0"
+version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
+checksum = "f1fdabc7756949593fe60f30ec81974b613357de856987752631dea1e3394c80"

 [[package]]
 name = "backtrace"
-version = "0.3.70"
+version = "0.3.71"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "95d8e92cac0961e91dbd517496b00f7e9b92363dbe6d42c3198268323798860c"
+checksum = "26b05800d2e817c8b3b4b54abd461726265fa9789ae34330622f2db9ee696f9d"
 dependencies = [
 "addr2line",
 "cc",
@@ -716,9 +716,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"

 [[package]]
 name = "bytes"
-version = "1.5.0"
+version = "1.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a2bd12c1caf447e69cd4528f47f94d203fd2582878ecb9e9465484c4148a8223"
+checksum = "514de17de45fdb8dc022b1a7975556c53c86f9f0aa5f534b98977b171857c2c9"
 dependencies = [
 "serde",
 ]
@@ -734,9 +734,9 @@ dependencies = [

 [[package]]
 name = "cargo-platform"
-version = "0.1.7"
+version = "0.1.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "694c8807f2ae16faecc43dc17d74b3eb042482789fd0eb64b39a2e04e087053f"
+checksum = "24b1f0365a6c6bb4020cd05806fd0d33c44d38046b8bd7f0e40814b9763cabfc"
 dependencies = [
 "serde",
 ]
@@ -779,9 +779,9 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"

 [[package]]
 name = "chrono"
-version = "0.4.35"
+version = "0.4.37"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8eaf5903dcbc0a39312feb77df2ff4c76387d591b9fc7b04a238dcf8bb62639a"
+checksum = "8a0d04d43504c61aa6c7531f1871dd0d418d91130162063b789da00fd7057a5e"
 dependencies = [
 "android-tzdata",
 "iana-time-zone",
@@ -814,9 +814,9 @@ dependencies = [

 [[package]]
 name = "clap"
-version = "4.5.3"
+version = "4.5.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "949626d00e063efc93b6dca932419ceb5432f99769911c0b995f7e884c778813"
+checksum = "90bc066a67923782aa8515dbaea16946c5bcc5addbd668bb80af688e53e548a0"
 dependencies = [
 "clap_builder",
 "clap_derive",
@@ -836,14 +836,14 @@ dependencies = [

 [[package]]
 name = "clap_derive"
-version = "4.5.3"
+version = "4.5.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "90239a040c80f5e14809ca132ddc4176ab33d5e17e49691793296e3fcb34d72f"
+checksum = "528131438037fd55894f62d6e9f068b8f45ac57ffa77517819645d10aed04f64"
 dependencies = [
 "heck 0.5.0",
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -1177,9 +1177,9 @@ dependencies = [

 [[package]]
 name = "der"
-version = "0.7.8"
+version = "0.7.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fffa369a668c8af7dbf8b5e56c9f744fbd399949ed171606040001947de40b1c"
+checksum = "f55bf8e7b65898637379c1b74eb1551107c8294ed26d855ceb9fd1a09cfc9bc0"
 dependencies = [
 "const-oid",
 "zeroize",
@@ -1409,7 +1409,7 @@ checksum = "6fd000fd6988e73bbe993ea3db9b1aa64906ab88766d654973924340c8cddb42"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -1575,7 +1575,7 @@ dependencies = [
 "regex",
 "serde",
 "serde_json",
- "syn 2.0.53",
+ "syn 2.0.58",
 "toml",
 "walkdir",
 ]
@@ -1593,7 +1593,7 @@ dependencies = [
 "proc-macro2",
 "quote",
 "serde_json",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -1619,7 +1619,7 @@ dependencies = [
 "serde",
 "serde_json",
 "strum",
- "syn 2.0.53",
+ "syn 2.0.58",
 "tempfile",
 "thiserror",
 "tiny-keccak",
@@ -1772,7 +1772,7 @@ dependencies = [
 "ark-std 0.3.0",
 "bincode",
 "chrono",
- "clap 4.5.3",
+ "clap 4.5.4",
 "colored",
 "colored_json",
 "console_error_panic_hook",
@@ -1837,9 +1837,9 @@ checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7"

 [[package]]
 name = "fastrand"
-version = "2.0.1"
+version = "2.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "25cbce373ec4653f1a01a31e8a5e5ec0c622dc27ff9c4e6606eefef5cbbed4a5"
+checksum = "658bd65b1cf4c852a3cc96f18a8ce7b5640f6b703f905c7d74532294c2a63984"

 [[package]]
 name = "fastrlp"
@@ -2028,7 +2028,7 @@ checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -2258,7 +2258,7 @@ dependencies = [
 [[package]]
 name = "halo2curves"
 version = "0.3.2"
-source = "git+https://github.com/privacy-scaling-explorations/halo2curves?tag=0.3.2#9f5c50810bbefe779ee5cf1d852b2fe85dc35d5e"
+source = "git+https://github.com/privacy-scaling-explorations/halo2curves.git?tag=0.3.2#9f5c50810bbefe779ee5cf1d852b2fe85dc35d5e"
 dependencies = [
 "ff",
 "group",
@@ -2602,9 +2602,9 @@ checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683"

 [[package]]
 name = "indexmap"
-version = "2.2.5"
+version = "2.2.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7b0b929d511467233429c45a44ac1dcaa21ba0f5ba11e4879e6ed28ddb4f9df4"
+checksum = "168fb715dda47215e360912c096649d23d58bf392ac62f73919e831745e40f26"
 dependencies = [
 "equivalent",
 "hashbrown 0.14.3",
@@ -2626,9 +2626,9 @@ dependencies = [

 [[package]]
 name = "indoc"
-version = "2.0.4"
+version = "2.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1e186cfbae8084e513daff4240b4797e342f988cecda4fb6c939150f96315fd8"
+checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5"

 [[package]]
 name = "inout"
@@ -2700,10 +2700,19 @@ dependencies = [
 ]

 [[package]]
-name = "itoa"
-version = "1.0.10"
+name = "itertools"
+version = "0.12.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b1a46d1a171d865aa5f83f92695765caa047a9b4cbae2cbf37dbd613a793fd4c"
+checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569"
+dependencies = [
+ "either",
+]
+
+[[package]]
+name = "itoa"
+version = "1.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b"

 [[package]]
 name = "jobserver"
@@ -2823,13 +2832,12 @@ checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058"

 [[package]]
 name = "libredox"
-version = "0.0.1"
+version = "0.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "85c833ca1e66078851dba29046874e38f08b2c883700aa29a03ddd3b23814ee8"
+checksum = "c0ff37bd590ca25063e35af745c343cb7a0271906fb7b37e4813e8f79f00268d"
 dependencies = [
 "bitflags 2.5.0",
 "libc",
- "redox_syscall",
 ]

 [[package]]
@@ -2877,7 +2885,7 @@ checksum = "fc2fb41a9bb4257a3803154bdf7e2df7d45197d1941c9b1a90ad815231630721"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -2962,24 +2970,24 @@ dependencies = [

 [[package]]
 name = "memchr"
-version = "2.7.1"
+version = "2.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149"
+checksum = "6c8640c5d730cb13ebd907d8d04b52f55ac9a2eec55b440c8892f40d56c76c1d"

 [[package]]
 name = "memmap2"
-version = "0.5.10"
+version = "0.9.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "83faa42c0a078c393f6b29d5db232d8be22776a891f8f56e5284faee4a20b327"
+checksum = "fe751422e4a8caa417e13c3ea66452215d7d63e19e604f4980461212f3ae1322"
 dependencies = [
 "libc",
 ]

 [[package]]
 name = "memoffset"
-version = "0.9.0"
+version = "0.9.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
 dependencies = [
 "autocfg",
 ]
@@ -3187,7 +3195,7 @@ dependencies = [
 "proc-macro-crate 3.1.0",
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -3271,7 +3279,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -3291,9 +3299,9 @@ dependencies = [

 [[package]]
 name = "openssl-sys"
-version = "0.9.101"
+version = "0.9.102"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dda2b0f344e78efc2facf7d195d098df0dd72151b26ab98da807afc26c198dff"
+checksum = "c597637d56fbc83893a35eb0dd04b2b8e7a50c91e64e9493e398b5df4fb45fa2"
 dependencies = [
 "cc",
 "libc",
@@ -3444,9 +3452,9 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e"

 [[package]]
 name = "pest"
-version = "2.7.8"
+version = "2.7.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "56f8023d0fb78c8e03784ea1c7f3fa36e68a723138990b8d5a47d916b651e7a8"
+checksum = "311fb059dee1a7b802f036316d790138c613a4e8b180c822e3925a662e9f0c95"
 dependencies = [
 "memchr",
 "thiserror",
@@ -3455,9 +3463,9 @@ dependencies = [

 [[package]]
 name = "pest_derive"
-version = "2.7.8"
+version = "2.7.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b0d24f72393fd16ab6ac5738bc33cdb6a9aa73f8b902e8fe29cf4e67d7dd1026"
+checksum = "f73541b156d32197eecda1a4014d7f868fd2bcb3c550d5386087cfba442bf69c"
 dependencies = [
 "pest",
 "pest_generator",
@@ -3465,22 +3473,22 @@ dependencies = [

 [[package]]
 name = "pest_generator"
-version = "2.7.8"
+version = "2.7.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fdc17e2a6c7d0a492f0158d7a4bd66cc17280308bbaff78d5bef566dca35ab80"
+checksum = "c35eeed0a3fab112f75165fdc026b3913f4183133f19b49be773ac9ea966e8bd"
 dependencies = [
 "pest",
 "pest_meta",
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
 name = "pest_meta"
-version = "2.7.8"
+version = "2.7.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "934cd7631c050f4674352a6e835d5f6711ffbfb9345c2fc0107155ac495ae293"
+checksum = "2adbf29bb9776f28caece835398781ab24435585fe0d4dc1374a61db5accedca"
 dependencies = [
 "once_cell",
 "pest",
@@ -3550,7 +3558,7 @@ dependencies = [
 "phf_shared 0.11.2",
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -3588,14 +3596,14 @@ checksum = "2f38a4412a78282e09a2cf38d195ea5420d15ba0602cb375210efbc877243965"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
 name = "pin-project-lite"
-version = "0.2.13"
+version = "0.2.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8afb450f006bf6385ca15ef45d71d2288452bc3683ce2e2cacc0d18e4be60b58"
+checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02"

 [[package]]
 name = "pin-utils"
@@ -3719,12 +3727,12 @@ checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c"

 [[package]]
 name = "prettyplease"
-version = "0.2.16"
+version = "0.2.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a41cf62165e97c7f814d2221421dbb9afcbcdb0a88068e5ea206e19951c2cbb5"
+checksum = "8d3928fb5db768cb86f891ff014f0144589297e3c6a1aba6ed7cecfdace270c7"
 dependencies = [
 "proc-macro2",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -3933,7 +3941,7 @@ dependencies = [
 "proc-macro2",
 "pyo3-macros-backend",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -3946,7 +3954,7 @@ dependencies = [
 "proc-macro2",
 "pyo3-build-config",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -4049,9 +4057,9 @@ checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"

 [[package]]
 name = "rayon"
-version = "1.9.0"
+version = "1.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e4963ed1bc86e4f3ee217022bd855b297cef07fb9eac5dfa1f788b220b49b3bd"
+checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa"
 dependencies = [
 "either",
 "rayon-core",
@@ -4089,9 +4097,9 @@ dependencies = [

 [[package]]
 name = "redox_users"
-version = "0.4.4"
+version = "0.4.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a18479200779601e498ada4e8c1e1f50e3ee19deb0259c25825a98b5603b2cb4"
+checksum = "bd283d9651eeda4b2a83a43c1c91b266c40fd76ecd39a50a8c630ae69dc72891"
 dependencies = [
 "getrandom",
 "libredox",
@@ -4100,9 +4108,9 @@ dependencies = [

 [[package]]
 name = "regex"
-version = "1.10.3"
+version = "1.10.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b62dbe01f0b06f9d8dc7d49e05a0785f153b00b2c227856282f671e0318c9b15"
+checksum = "c117dbdfde9c8308975b6a18d71f3f385c89461f7b3fb054288ecf2a2058ba4c"
 dependencies = [
 "aho-corasick",
 "memchr",
@@ -4123,9 +4131,9 @@ dependencies = [

 [[package]]
 name = "regex-syntax"
-version = "0.8.2"
+version = "0.8.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f"
+checksum = "adad44e29e4c806119491a7f06f03de4d1af22c3a680dd47f1e6e179439d1f56"

 [[package]]
 name = "remove_dir_all"
@@ -4444,9 +4452,9 @@ dependencies = [

 [[package]]
 name = "scale-info"
-version = "2.11.0"
+version = "2.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2ef2175c2907e7c8bc0a9c3f86aeb5ec1f3b275300ad58a44d0c3ae379a5e52e"
+checksum = "788745a868b0e751750388f4e6546eb921ef714a4317fa6954f7cde114eb2eb7"
 dependencies = [
 "cfg-if",
 "derive_more",
@@ -4456,9 +4464,9 @@ dependencies = [

 [[package]]
 name = "scale-info-derive"
-version = "2.11.0"
+version = "2.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "634d9b8eb8fd61c5cdd3390d9b2132300a7e7618955b98b8416f118c1b4e144f"
+checksum = "7dc2f4e8bc344b9fc3d5f74f72c2e55bfc38d28dc2ebc69c194a3df424e4d9ac"
 dependencies = [
 "proc-macro-crate 1.3.1",
 "proc-macro2",
@@ -4524,9 +4532,9 @@ dependencies = [

 [[package]]
 name = "security-framework"
-version = "2.9.2"
+version = "2.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "05b64fb303737d99b81884b2c63433e9ae28abebe5eb5045dcdd175dc2ecf4de"
+checksum = "770452e37cad93e0a50d5abc3990d2bc351c36d0328f86cefec2f2fb206eaef6"
 dependencies = [
 "bitflags 1.3.2",
 "core-foundation",
@@ -4537,9 +4545,9 @@ dependencies = [

 [[package]]
 name = "security-framework-sys"
-version = "2.9.1"
+version = "2.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e932934257d3b408ed8f30db49d85ea163bfe74961f017f405b025af298f0c7a"
+checksum = "41f3cc463c0ef97e11c3461a9d3787412d30e8e7eb907c79180c4a57bf7c04ef"
 dependencies = [
 "core-foundation-sys",
 "libc",
@@ -4601,9 +4609,9 @@ dependencies = [

 [[package]]
 name = "serde-wasm-bindgen"
-version = "0.4.5"
+version = "0.6.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e3b4c031cd0d9014307d82b8abf653c0290fbdaeb4c02d00c63cf52f728628bf"
+checksum = "8302e169f0eddcc139c70f139d19d6467353af16f9fce27e8c30158036a1e16b"
 dependencies = [
 "js-sys",
 "serde",
@@ -4637,14 +4645,14 @@ checksum = "7eb0b34b42edc17f6b7cac84a52a1c5f0e1bb2227e997ca9011ea3dd34e8610b"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
 name = "serde_json"
-version = "1.0.114"
+version = "1.0.115"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c5f09b1bd632ef549eaa9f60a1f8de742bdbc698e6cee2095fc84dde5f549ae0"
+checksum = "12dc5c46daa8e9fdf4f5e71b6cf9a53f2487da0e86e55808e2d35539666497dd"
 dependencies = [
 "itoa",
 "ryu",
@@ -4845,12 +4853,12 @@ checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82"

 [[package]]
 name = "string-interner"
-version = "0.14.0"
+version = "0.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "91e2531d8525b29b514d25e275a43581320d587b86db302b9a7e464bac579648"
+checksum = "07f9fdfdd31a0ff38b59deb401be81b73913d76c9cc5b1aed4e1330a223420b9"
 dependencies = [
 "cfg-if",
- "hashbrown 0.11.2",
+ "hashbrown 0.14.3",
 "serde",
 ]

@@ -4880,9 +4888,9 @@ dependencies = [

 [[package]]
 name = "strsim"
-version = "0.11.0"
+version = "0.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5ee073c9e4cd00e28217186dbe12796d692868f432bf2e97ee73bed0c56dfa01"
+checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"

 [[package]]
 name = "strum"
@@ -4903,7 +4911,7 @@ dependencies = [
 "proc-macro2",
 "quote",
 "rustversion",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -4938,9 +4946,9 @@ dependencies = [

 [[package]]
 name = "syn"
-version = "2.0.53"
+version = "2.0.58"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7383cd0e49fff4b6b90ca5670bfd3e9d6a733b3f90c686605aa7eec8c4996032"
+checksum = "44cfb93f38070beee36b3fef7d4f5a16f27751d94b187b666a5cc5e9b0d30687"
 dependencies = [
 "proc-macro2",
 "quote",
@@ -5113,7 +5121,7 @@ checksum = "c61f3ba182994efc43764a46c018c347bc492c79f024e705f46567b418f6d4f7"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -5179,9 +5187,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"

 [[package]]
 name = "tokio"
-version = "1.36.0"
+version = "1.37.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "61285f6515fa018fb2d1e46eb21223fff441ee8db5d0f1435e8ab4f5cdb80931"
+checksum = "1adbebffeca75fcfd058afa480fb6c0b81e165a0323f9c9d39c9697e37c46787"
 dependencies = [
 "backtrace",
 "bytes",
@@ -5202,7 +5210,7 @@ checksum = "5b8a1e28f2deaa14e508979454cb3a223b10b938b45af148bc0986de36f1923b"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -5337,7 +5345,7 @@ source = "git+https://github.com/zkonduit/enum_to_subcommand#42e9870f1f757932bab
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -5365,7 +5373,7 @@ checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -5389,8 +5397,8 @@ dependencies = [

 [[package]]
 name = "tract-core"
-version = "0.20.23-pre"
-source = "git+https://github.com/sonos/tract/?rev=7b1aa33b2f7d1f19b80e270c83320f0f94daff69#7b1aa33b2f7d1f19b80e270c83320f0f94daff69"
+version = "0.21.3"
+source = "git+https://github.com/sonos/tract/?rev=681a096f02c9d7d363102d9fb0e446d1710ac2c8#681a096f02c9d7d363102d9fb0e446d1710ac2c8"
 dependencies = [
 "anyhow",
 "bit-set",
@@ -5413,12 +5421,12 @@ dependencies = [

 [[package]]
 name = "tract-data"
-version = "0.20.23-pre"
-source = "git+https://github.com/sonos/tract/?rev=7b1aa33b2f7d1f19b80e270c83320f0f94daff69#7b1aa33b2f7d1f19b80e270c83320f0f94daff69"
+version = "0.21.3"
+source = "git+https://github.com/sonos/tract/?rev=681a096f02c9d7d363102d9fb0e446d1710ac2c8#681a096f02c9d7d363102d9fb0e446d1710ac2c8"
 dependencies = [
 "anyhow",
 "half 2.2.1",
- "itertools 0.10.5",
+ "itertools 0.12.1",
 "lazy_static",
 "maplit",
 "ndarray",
@@ -5432,8 +5440,8 @@ dependencies = [

 [[package]]
 name = "tract-hir"
-version = "0.20.23-pre"
-source = "git+https://github.com/sonos/tract/?rev=7b1aa33b2f7d1f19b80e270c83320f0f94daff69#7b1aa33b2f7d1f19b80e270c83320f0f94daff69"
+version = "0.21.3"
+source = "git+https://github.com/sonos/tract/?rev=681a096f02c9d7d363102d9fb0e446d1710ac2c8#681a096f02c9d7d363102d9fb0e446d1710ac2c8"
 dependencies = [
 "derive-new",
 "log",
@@ -5442,8 +5450,8 @@ dependencies = [

 [[package]]
 name = "tract-linalg"
-version = "0.20.23-pre"
-source = "git+https://github.com/sonos/tract/?rev=7b1aa33b2f7d1f19b80e270c83320f0f94daff69#7b1aa33b2f7d1f19b80e270c83320f0f94daff69"
+version = "0.21.3"
+source = "git+https://github.com/sonos/tract/?rev=681a096f02c9d7d363102d9fb0e446d1710ac2c8#681a096f02c9d7d363102d9fb0e446d1710ac2c8"
 dependencies = [
 "cc",
 "derive-new",
@@ -5466,8 +5474,8 @@ dependencies = [

 [[package]]
 name = "tract-nnef"
-version = "0.20.23-pre"
-source = "git+https://github.com/sonos/tract/?rev=7b1aa33b2f7d1f19b80e270c83320f0f94daff69#7b1aa33b2f7d1f19b80e270c83320f0f94daff69"
+version = "0.21.3"
+source = "git+https://github.com/sonos/tract/?rev=681a096f02c9d7d363102d9fb0e446d1710ac2c8#681a096f02c9d7d363102d9fb0e446d1710ac2c8"
 dependencies = [
 "byteorder",
 "flate2",
@@ -5480,8 +5488,8 @@ dependencies = [

 [[package]]
 name = "tract-onnx"
-version = "0.20.23-pre"
-source = "git+https://github.com/sonos/tract/?rev=7b1aa33b2f7d1f19b80e270c83320f0f94daff69#7b1aa33b2f7d1f19b80e270c83320f0f94daff69"
+version = "0.21.3"
+source = "git+https://github.com/sonos/tract/?rev=681a096f02c9d7d363102d9fb0e446d1710ac2c8#681a096f02c9d7d363102d9fb0e446d1710ac2c8"
 dependencies = [
 "bytes",
 "derive-new",
@@ -5497,8 +5505,8 @@ dependencies = [

 [[package]]
 name = "tract-onnx-opl"
-version = "0.20.23-pre"
-source = "git+https://github.com/sonos/tract/?rev=7b1aa33b2f7d1f19b80e270c83320f0f94daff69#7b1aa33b2f7d1f19b80e270c83320f0f94daff69"
+version = "0.21.3"
+source = "git+https://github.com/sonos/tract/?rev=681a096f02c9d7d363102d9fb0e446d1710ac2c8#681a096f02c9d7d363102d9fb0e446d1710ac2c8"
 dependencies = [
 "getrandom",
 "log",
@@ -5755,7 +5763,7 @@ dependencies = [
 "once_cell",
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 "wasm-bindgen-shared",
 ]

@@ -5799,7 +5807,7 @@ checksum = "e94f17b526d0a461a191c78ea52bbce64071ed5c04c9ffe424dcb38f74171bb7"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 "wasm-bindgen-backend",
 "wasm-bindgen-shared",
 ]
@@ -5844,7 +5852,7 @@ checksum = "b7f89739351a2e03cb94beb799d47fb2cac01759b40ec441f7de39b00cbf7ef0"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -6154,7 +6162,7 @@ checksum = "9ce1b18ccd8e73a9321186f97e46f9f04b778851177567b1975109d26a08d2a6"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]

 [[package]]
@@ -6174,5 +6182,5 @@ checksum = "ce36e65b0d2999d2aafac989fb249189a141aee1f53c612c1f37d72631959f69"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn 2.0.53",
+ "syn 2.0.58",
 ]
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -80,7 +80,7 @@ pyo3-asyncio = { version = "0.20.0", features = [
    "tokio-runtime",
 ], default_features = false, optional = true }
 pyo3-log = { version = "0.9.0", default_features = false, optional = true }
-tract-onnx = { git = "https://github.com/sonos/tract/", rev = "7b1aa33b2f7d1f19b80e270c83320f0f94daff69", default_features = false, optional = true }
+tract-onnx = { git = "https://github.com/sonos/tract/", rev = "681a096f02c9d7d363102d9fb0e446d1710ac2c8", default_features = false, optional = true }
 tabled = { version = "0.12.0", optional = true }


@@ -95,10 +95,10 @@ getrandom = { version = "0.2.8", features = ["js"] }
 instant = { version = "0.1", features = ["wasm-bindgen", "inaccurate"] }

 [target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies]
-wasm-bindgen-rayon = { version = "1.0", optional = true }
-wasm-bindgen-test = "0.3.34"
-serde-wasm-bindgen = "0.4"
-wasm-bindgen = { version = "0.2.81", features = ["serde-serialize"] }
+wasm-bindgen-rayon = { version = "1.2.1", optional = true }
+wasm-bindgen-test = "0.3.42"
+serde-wasm-bindgen = "0.6.5"
+wasm-bindgen = { version = "0.2.92", features = ["serde-serialize"] }
 console_error_panic_hook = "0.1.7"
 wasm-bindgen-console-logger = "0.1.1"

--- a/benches/poseidon.rs
+++ b/benches/poseidon.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
 use ezkl::circuit::modules::poseidon::spec::{PoseidonSpec, POSEIDON_RATE, POSEIDON_WIDTH};
 use ezkl::circuit::modules::poseidon::{PoseidonChip, PoseidonConfig};
@@ -48,7 +50,7 @@ impl Circuit<Fr> for MyCircuit {
    ) -> Result<(), Error> {
        let chip: PoseidonChip<PoseidonSpec, POSEIDON_WIDTH, POSEIDON_RATE, L> =
            PoseidonChip::new(config);
-        chip.layout(&mut layouter, &[self.image.clone()], 0)?;
+        chip.layout(&mut layouter, &[self.image.clone()], 0, &mut HashMap::new())?;
        Ok(())
    }
 }
--- a/examples/notebooks/keras_simple_demo.ipynb
+++ b/examples/notebooks/keras_simple_demo.ipynb
@@ -67,6 +67,7 @@
                "model.add(Dense(128, activation='relu'))\n",
                "model.add(Dropout(0.5))\n",
                "model.add(Dense(10, activation='softmax'))\n",
+                "model.output_names=['output']\n",
                "\n",
                "\n",
                "# Train the model as you like here (skipped for brevity)\n",
--- a/examples/notebooks/mnist_gan.ipynb
+++ b/examples/notebooks/mnist_gan.ipynb
@@ -38,7 +38,7 @@
                "import logging\n",
                "\n",
                "import tensorflow as tf\n",
-                "from tensorflow.keras.optimizers.legacy import Adam\n",
+                "from tensorflow.keras.optimizers import Adam\n",
                "from tensorflow.keras.layers import *\n",
                "from tensorflow.keras.models import Model\n",
                "from tensorflow.keras.datasets import mnist\n",
@@ -71,9 +71,11 @@
            },
            "outputs": [],
            "source": [
-                "opt = Adam()\n",
                "ZDIM = 100\n",
                "\n",
+                "opt = Adam()\n",
+                "\n",
+                "\n",
                "# discriminator\n",
                "# 0 if it's fake, 1 if it's real\n",
                "x = in1 = Input((28,28))\n",
@@ -114,8 +116,11 @@
                "\n",
                "gm = Model(in1, x)\n",
                "gm.compile('adam', 'mse')\n",
+                "gm.output_names=['output']\n",
                "gm.summary()\n",
                "\n",
+                "opt = Adam()\n",
+                "\n",
                "# GAN\n",
                "dm.trainable = False\n",
                "x = dm(gm.output)\n",
@@ -415,7 +420,7 @@
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
-            "version": "3.9.15"
+            "version": "3.12.2"
        }
    },
    "nbformat": 4,
--- a/examples/notebooks/mnist_gan_proof_splitting.ipynb
+++ b/examples/notebooks/mnist_gan_proof_splitting.ipynb
--- a/examples/notebooks/mnist_vae.ipynb
+++ b/examples/notebooks/mnist_vae.ipynb
@@ -349,6 +349,8 @@
                "z_log_var = Dense(ZDIM)(x)\n",
                "z = Lambda(lambda x: x[0] + K.exp(0.5 * x[1]) * K.random_normal(shape=K.shape(x[0])))([z_mu, z_log_var])\n",
                "dec = get_decoder()\n",
+                "dec.output_names=['output']\n",
+                "\n",
                "out = dec(z)\n",
                "\n",
                "mse_loss = mse(Reshape((28*28,))(in1), Reshape((28*28,))(out)) * 28 * 28\n",
--- a/examples/notebooks/random_forest.ipynb
+++ b/examples/notebooks/random_forest.ipynb
@@ -61,11 +61,10 @@
                "from sklearn.datasets import load_iris\n",
                "from sklearn.model_selection import train_test_split\n",
                "from sklearn.ensemble import RandomForestClassifier as Rf\n",
-                "import sk2torch\n",
                "import torch\n",
                "import ezkl\n",
                "import os\n",
-                "from torch import nn\n",
+                "from hummingbird.ml import convert\n",
                "\n",
                "\n",
                "\n",
@@ -77,28 +76,12 @@
                "clr.fit(X_train, y_train)\n",
                "\n",
                "\n",
-                "trees = []\n",
-                "for tree in clr.estimators_:\n",
-                "    trees.append(sk2torch.wrap(tree))\n",
                "\n",
-                "\n",
-                "class RandomForest(nn.Module):\n",
-                "    def __init__(self, trees):\n",
-                "        super(RandomForest, self).__init__()\n",
-                "        self.trees = nn.ModuleList(trees)\n",
-                "\n",
-                "    def forward(self, x):\n",
-                "        out = self.trees[0](x)\n",
-                "        for tree in self.trees[1:]:\n",
-                "            out += tree(x)\n",
-                "        return out / len(self.trees)\n",
-                "\n",
-                "\n",
-                "torch_rf = RandomForest(trees)\n",
+                "torch_rf = convert(clr, 'torch')\n",
                "# assert predictions from torch are = to sklearn \n",
                "diffs = []\n",
                "for i in range(len(X_test)):\n",
-                "    torch_pred = torch_rf(torch.tensor(X_test[i].reshape(1, -1)))\n",
+                "    torch_pred = torch_rf.predict(torch.tensor(X_test[i].reshape(1, -1)))\n",
                "    sk_pred = clr.predict(X_test[i].reshape(1, -1))\n",
                "    diffs.append(torch_pred[0].round() - sk_pred[0])\n",
                "\n",
@@ -134,14 +117,12 @@
                "\n",
                "# export to onnx format\n",
                "\n",
-                "torch_rf.eval()\n",
-                "\n",
                "# Input to the model\n",
                "shape = X_train.shape[1:]\n",
                "x = torch.rand(1, *shape, requires_grad=False)\n",
-                "torch_out = torch_rf(x)\n",
+                "torch_out = torch_rf.predict(x)\n",
                "# Export the model\n",
-                "torch.onnx.export(torch_rf,               # model being run\n",
+                "torch.onnx.export(torch_rf.model,               # model being run\n",
                "                  # model input (or a tuple for multiple inputs)\n",
                "                  x,\n",
                "                  # where to save the model (can be a file or file-like object)\n",
@@ -158,7 +139,7 @@
                "\n",
                "data = dict(input_shapes=[shape],\n",
                "            input_data=[d],\n",
-                "            output_data=[((o).detach().numpy()).reshape([-1]).tolist() for o in torch_out])\n",
+                "            output_data=[o.reshape([-1]).tolist() for o in torch_out])\n",
                "\n",
                "# Serialize data into file:\n",
                "json.dump(data, open(\"input.json\", 'w'))\n"
@@ -321,7 +302,7 @@
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
-            "version": "3.9.15"
+            "version": "3.12.2"
        }
    },
    "nbformat": 4,
--- a/examples/notebooks/tictactoe_autoencoder.ipynb
+++ b/examples/notebooks/tictactoe_autoencoder.ipynb
@@ -13,7 +13,7 @@
        },
        {
            "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 1,
            "metadata": {},
            "outputs": [],
            "source": [
@@ -57,7 +57,7 @@
        },
        {
            "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 2,
            "metadata": {},
            "outputs": [],
            "source": [
@@ -119,7 +119,7 @@
        },
        {
            "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 3,
            "metadata": {},
            "outputs": [],
            "source": [
@@ -163,7 +163,7 @@
        },
        {
            "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 4,
            "metadata": {},
            "outputs": [],
            "source": [
@@ -217,7 +217,7 @@
        },
        {
            "cell_type": "code",
-            "execution_count": null,
+            "execution_count": 6,
            "metadata": {},
            "outputs": [],
            "source": [
@@ -227,6 +227,10 @@
                "        self.length = self.compute_length(self.file_good)\n",
                "        self.data = self.load_data(self.file_good)\n",
                "\n",
+                "    def __iter__(self):\n",
+                "        for i in range(len(self.data)):\n",
+                "            yield self.data[i]\n",
+                "\n",
                "    def parse_json_object(self, line):\n",
                "        try:\n",
                "            return json.loads(line)\n",
@@ -749,7 +753,7 @@
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
-            "version": "3.9.15"
+            "version": "3.12.2"
        }
    },
    "nbformat": 4,
--- a/examples/notebooks/tictactoe_binary_classification.ipynb
+++ b/examples/notebooks/tictactoe_binary_classification.ipynb
@@ -209,6 +209,11 @@
                "        self.length = self.compute_length(self.file_good, self.file_bad)\n",
                "        self.data = self.load_data(self.file_good, self.file_bad)\n",
                "\n",
+                "    def __iter__(self):\n",
+                "        for i in range(len(self.data)):\n",
+                "            yield self.data[i]\n",
+                "\n",
+                "\n",
                "    def parse_json_object(self, line):\n",
                "        try:\n",
                "            return json.loads(line)\n",
@@ -637,7 +642,7 @@
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
-            "version": "3.9.15"
+            "version": "3.12.2"
        }
    },
    "nbformat": 4,
--- a/examples/onnx/lstm_large/input.json
+++ b/examples/onnx/lstm_large/input.json
@@ -0,0 +1,13 @@
+{
+  "input_data": [
+    [
+      0.8894134163856506,
+      0.8894201517105103
+    ]
+  ],
+  "output_data": [
+    [
+      0.8436377
+    ]
+  ]
+}
--- a/examples/onnx/lstm_large/network.onnx
+++ b/examples/onnx/lstm_large/network.onnx
--- a/in-browser-evm-verifier/package.json
+++ b/in-browser-evm-verifier/package.json
@@ -17,7 +17,7 @@
    "clean": "rm -r dist || true",
    "build:commonjs": "tsc --project tsconfig.commonjs.json && resolve-tspaths -p tsconfig.commonjs.json",
    "build:esm": "tsc --project tsconfig.esm.json && resolve-tspaths -p tsconfig.esm.json",
-    "build": "pnpm run clean && pnpm run build:commonjs && pnpm run build:esm"
+    "build": "npm run clean && npm run build:commonjs && npm run build:esm"
  },
  "dependencies": {
    "@ethereumjs/common": "^4.0.0",
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,14 +1,14 @@
-attrs==22.2.0
-exceptiongroup==1.1.1
-importlib-metadata==6.1.0
+attrs==23.2.0
+exceptiongroup==1.2.0
+importlib-metadata==7.1.0
 iniconfig==2.0.0
-maturin==1.0.1
-packaging==23.0
-pluggy==1.0.0
-pytest==7.2.2
+maturin==1.5.0
+packaging==24.0
+pluggy==1.4.0
+pytest==8.1.1
 tomli==2.0.1
-typing-extensions==4.5.0
-zipp==3.15.0
-onnx==1.14.1
-onnxruntime==1.14.1
-numpy==1.21.6
+typing-extensions==4.10.0
+zipp==3.18.1
+onnx==1.15.0
+onnxruntime==1.17.1
+numpy==1.26.4
--- a/src/circuit/modules/mod.rs
+++ b/src/circuit/modules/mod.rs
@@ -15,6 +15,8 @@ pub use planner::*;

 use crate::tensor::{TensorType, ValTensor};

+use super::region::ConstantsMap;
+
 /// Module trait used to extend ezkl functionality
 pub trait Module<F: PrimeField + TensorType + PartialOrd> {
    /// Config
@@ -39,6 +41,7 @@ pub trait Module<F: PrimeField + TensorType + PartialOrd> {
        &self,
        layouter: &mut impl Layouter<F>,
        input: &[ValTensor<F>],
+        constants: &mut ConstantsMap<F>,
    ) -> Result<Self::InputAssignments, Error>;
    /// Layout
    fn layout(
@@ -46,6 +49,7 @@ pub trait Module<F: PrimeField + TensorType + PartialOrd> {
        layouter: &mut impl Layouter<F>,
        input: &[ValTensor<F>],
        row_offset: usize,
+        constants: &mut ConstantsMap<F>,
    ) -> Result<ValTensor<F>, Error>;
    /// Number of instance values the module uses every time it is applied
    fn instance_increment_input(&self) -> Vec<usize>;
--- a/src/circuit/modules/polycommit.rs
+++ b/src/circuit/modules/polycommit.rs
@@ -4,6 +4,8 @@ is already implemented in halo2_gadgets, there is no wrapper chip that makes it
 Thanks to https://github.com/summa-dev/summa-solvency/blob/master/src/chips/poseidon/hash.rs for the inspiration (and also helping us understand how to use this).
 */

+use std::collections::HashMap;
+
 // This chip adds a set of advice columns to the gadget Chip to store the inputs of the hash
 use halo2_proofs::halo2curves::bn256::Fr as Fp;
 use halo2_proofs::poly::commitment::{Blind, CommitmentScheme, Params};
@@ -13,6 +15,7 @@ use halo2curves::group::prime::PrimeCurveAffine;
 use halo2curves::group::Curve;
 use halo2curves::CurveAffine;

+use crate::circuit::region::ConstantsMap;
 use crate::tensor::{Tensor, ValTensor, ValType, VarTensor};

 use super::Module;
@@ -107,6 +110,7 @@ impl Module<Fp> for PolyCommitChip {
        &self,
        _: &mut impl Layouter<Fp>,
        _: &[ValTensor<Fp>],
+        _: &mut ConstantsMap<Fp>,
    ) -> Result<Self::InputAssignments, Error> {
        Ok(())
    }
@@ -119,11 +123,24 @@ impl Module<Fp> for PolyCommitChip {
        layouter: &mut impl Layouter<Fp>,
        input: &[ValTensor<Fp>],
        _: usize,
+        constants: &mut ConstantsMap<Fp>,
    ) -> Result<ValTensor<Fp>, Error> {
        assert_eq!(input.len(), 1);
+
+        let local_constants = constants.clone();
        layouter.assign_region(
            || "PolyCommit",
-            |mut region| self.config.inputs.assign(&mut region, 0, &input[0]),
+            |mut region| {
+                let mut local_inner_constants = local_constants.clone();
+                let res = self.config.inputs.assign(
+                    &mut region,
+                    0,
+                    &input[0],
+                    &mut local_inner_constants,
+                )?;
+                *constants = local_inner_constants;
+                Ok(res)
+            },
        )
    }

@@ -184,7 +201,12 @@ mod tests {
            mut layouter: impl Layouter<Fp>,
        ) -> Result<(), Error> {
            let polycommit_chip = PolyCommitChip::new(config);
-            polycommit_chip.layout(&mut layouter, &[self.message.clone()], 0);
+            polycommit_chip.layout(
+                &mut layouter,
+                &[self.message.clone()],
+                0,
+                &mut HashMap::new(),
+            );

            Ok(())
        }
--- a/src/circuit/modules/poseidon.rs
+++ b/src/circuit/modules/poseidon.rs
@@ -18,6 +18,7 @@ use maybe_rayon::slice::ParallelSlice;

 use std::marker::PhantomData;

+use crate::circuit::region::ConstantsMap;
 use crate::tensor::{Tensor, ValTensor, ValType};

 use super::Module;
@@ -172,12 +173,15 @@ impl<S: Spec<Fp, WIDTH, RATE> + Sync, const WIDTH: usize, const RATE: usize, con
        &self,
        layouter: &mut impl Layouter<Fp>,
        message: &[ValTensor<Fp>],
+        constants: &mut ConstantsMap<Fp>,
    ) -> Result<Self::InputAssignments, Error> {
        assert_eq!(message.len(), 1);
        let message = message[0].clone();

        let start_time = instant::Instant::now();

+        let local_constants = constants.clone();
+
        let res = layouter.assign_region(
            || "load message",
            |mut region| {
@@ -199,12 +203,26 @@ impl<S: Spec<Fp, WIDTH, RATE> + Sync, const WIDTH: usize, const RATE: usize, con
                                ValType::PrevAssigned(v) | ValType::AssignedConstant(v, ..) => {
                                    Ok(v.clone())
                                }
-                                ValType::Constant(f) => region.assign_advice_from_constant(
-                                    || format!("load message_{}", i),
-                                    self.config.hash_inputs[x],
-                                    y,
-                                    *f,
-                                ),
+                                ValType::Constant(f) => {
+                                    if local_constants.contains_key(f) {
+                                        Ok(constants.get(f).unwrap().assigned_cell().ok_or({
+                                            log::error!("constant not previously assigned");
+                                            Error::Synthesis
+                                        })?)
+                                    } else {
+                                        let res = region.assign_advice_from_constant(
+                                            || format!("load message_{}", i),
+                                            self.config.hash_inputs[x],
+                                            y,
+                                            *f,
+                                        )?;
+
+                                        constants
+                                            .insert(*f, ValType::AssignedConstant(res.clone(), *f));
+
+                                        Ok(res)
+                                    }
+                                }
                                e => {
                                    log::error!(
                                        "wrong input type {:?}, must be previously assigned",
@@ -270,8 +288,9 @@ impl<S: Spec<Fp, WIDTH, RATE> + Sync, const WIDTH: usize, const RATE: usize, con
        layouter: &mut impl Layouter<Fp>,
        input: &[ValTensor<Fp>],
        row_offset: usize,
+        constants: &mut ConstantsMap<Fp>,
    ) -> Result<ValTensor<Fp>, Error> {
-        let (mut input_cells, zero_val) = self.layout_inputs(layouter, input)?;
+        let (mut input_cells, zero_val) = self.layout_inputs(layouter, input, constants)?;
        // extract the values from the input cells
        let mut assigned_input: Tensor<ValType<Fp>> =
            input_cells.iter().map(|e| ValType::from(e.clone())).into();
@@ -434,7 +453,7 @@ mod tests {
        *,
    };

-    use std::marker::PhantomData;
+    use std::{collections::HashMap, marker::PhantomData};

    use halo2_gadgets::poseidon::primitives::Spec;
    use halo2_proofs::{
@@ -477,7 +496,12 @@ mod tests {
            mut layouter: impl Layouter<Fp>,
        ) -> Result<(), Error> {
            let chip: PoseidonChip<PoseidonSpec, WIDTH, RATE, L> = PoseidonChip::new(config);
-            chip.layout(&mut layouter, &[self.message.clone()], 0)?;
+            chip.layout(
+                &mut layouter,
+                &[self.message.clone()],
+                0,
+                &mut HashMap::new(),
+            )?;

            Ok(())
        }
--- a/src/circuit/ops/chip.rs
+++ b/src/circuit/ops/chip.rs
@@ -345,7 +345,7 @@ pub struct BaseConfig<F: PrimeField + TensorType + PartialOrd> {
    _marker: PhantomData<F>,
 }

-impl<F: PrimeField + TensorType + PartialOrd> BaseConfig<F> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> BaseConfig<F> {
    /// Returns a new [BaseConfig] with no inputs, no selectors, and no tables.
    pub fn dummy(col_size: usize, num_inner_cols: usize) -> Self {
        Self {
--- a/src/circuit/ops/hybrid.rs
+++ b/src/circuit/ops/hybrid.rs
@@ -46,7 +46,8 @@ pub enum HybridOp {
        dim: usize,
    },
    Softmax {
-        scale: utils::F32,
+        input_scale: utils::F32,
+        output_scale: utils::F32,
        axes: Vec<usize>,
    },
    RangeCheck(Tolerance),
@@ -70,7 +71,7 @@ pub enum HybridOp {
    },
 }

-impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for HybridOp {
    ///
    fn requires_homogenous_input_scales(&self) -> Vec<usize> {
        match self {
@@ -130,9 +131,16 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
                kernel_shape,
                normalized,
            } => tensor::ops::sumpool(&x, *padding, *stride, *kernel_shape, *normalized)?,
-            HybridOp::Softmax { scale, axes } => {
-                tensor::ops::nonlinearities::softmax_axes(&x, scale.into(), axes)
-            }
+            HybridOp::Softmax {
+                input_scale,
+                output_scale,
+                axes,
+            } => tensor::ops::nonlinearities::softmax_axes(
+                &x,
+                input_scale.into(),
+                output_scale.into(),
+                axes,
+            ),
            HybridOp::RangeCheck(tol) => {
                let y = inputs[1].clone().map(|x| felt_to_i128(x));
                tensor::ops::nonlinearities::range_check_percent(&[x, y], 128, 128, tol.val)
@@ -203,8 +211,15 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
            ),
            HybridOp::ReduceMin { axes } => format!("REDUCEMIN (axes={:?})", axes),
            HybridOp::ReduceArgMin { dim } => format!("REDUCEARGMIN (dim={})", dim),
-            HybridOp::Softmax { scale, axes } => {
-                format!("SOFTMAX (scale={}, axes={:?})", scale, axes)
+            HybridOp::Softmax {
+                input_scale,
+                output_scale,
+                axes,
+            } => {
+                format!(
+                    "SOFTMAX (input_scale={}, output_scale={}, axes={:?})",
+                    input_scale, output_scale, axes
+                )
            }
            HybridOp::RangeCheck(p) => format!("RANGECHECK (tol={:?})", p),
            HybridOp::Greater => "GREATER".into(),
@@ -324,9 +339,18 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
            HybridOp::ReduceArgMin { dim } => {
                layouts::argmin_axes(config, region, values[..].try_into()?, *dim)?
            }
-            HybridOp::Softmax { scale, axes } => {
-                layouts::softmax_axes(config, region, values[..].try_into()?, *scale, axes)?
-            }
+            HybridOp::Softmax {
+                input_scale,
+                output_scale,
+                axes,
+            } => layouts::softmax_axes(
+                config,
+                region,
+                values[..].try_into()?,
+                *input_scale,
+                *output_scale,
+                axes,
+            )?,
            HybridOp::RangeCheck(tol) => layouts::range_check_percent(
                config,
                region,
@@ -359,8 +383,9 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for HybridOp {
            | HybridOp::ReduceArgMax { .. }
            | HybridOp::OneHot { .. }
            | HybridOp::ReduceArgMin { .. } => 0,
-            HybridOp::Softmax { .. } => 2 * in_scales[0],
-            HybridOp::Recip { output_scale, .. } => multiplier_to_scale(output_scale.0 as f64),
+            HybridOp::Softmax { output_scale, .. } | HybridOp::Recip { output_scale, .. } => {
+                multiplier_to_scale(output_scale.0 as f64)
+            }
            _ => in_scales[0],
        };
        Ok(scale)
--- a/src/circuit/ops/layouts.rs
+++ b/src/circuit/ops/layouts.rs
@@ -9,6 +9,7 @@ use halo2curves::ff::PrimeField;
 use itertools::Itertools;
 use log::{error, trace};
 use maybe_rayon::{
+    iter::IntoParallelRefIterator,
    prelude::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator},
    slice::ParallelSliceMut,
 };
@@ -33,7 +34,7 @@ use super::*;
 use crate::circuit::ops::lookup::LookupOp;

 /// Same as div but splits the division into N parts
-pub(crate) fn loop_div<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn loop_div<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    value: &[ValTensor<F>; 1],
@@ -68,7 +69,7 @@ pub(crate) fn loop_div<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Div accumulated layout
-pub(crate) fn div<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn div<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    value: &[ValTensor<F>; 1],
@@ -93,9 +94,9 @@ pub(crate) fn div<F: PrimeField + TensorType + PartialOrd>(
    let mut claimed_output: ValTensor<F> = if is_assigned {
        let input_evals = input.get_int_evals()?;
        tensor::ops::nonlinearities::const_div(&input_evals.clone(), felt_to_i128(div) as f64)
-            .iter()
-            .map(|x| Ok(Value::known(i128_to_felt(*x))))
-            .collect::<Result<Tensor<Value<F>>, Box<dyn Error>>>()?
+            .par_iter()
+            .map(|x| Value::known(i128_to_felt(*x)))
+            .collect::<Tensor<Value<F>>>()
            .into()
    } else {
        Tensor::new(
@@ -133,7 +134,7 @@ pub(crate) fn div<F: PrimeField + TensorType + PartialOrd>(
 }

 /// recip accumulated layout
-pub(crate) fn recip<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn recip<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    value: &[ValTensor<F>; 1],
@@ -166,9 +167,9 @@ pub(crate) fn recip<F: PrimeField + TensorType + PartialOrd>(
            felt_to_i128(input_scale) as f64,
            felt_to_i128(output_scale) as f64,
        )
-        .iter()
-        .map(|x| Ok(Value::known(i128_to_felt(*x))))
-        .collect::<Result<Tensor<Value<F>>, Box<dyn Error>>>()?
+        .par_iter()
+        .map(|x| Value::known(i128_to_felt(*x)))
+        .collect::<Tensor<Value<F>>>()
        .into()
    } else {
        Tensor::new(
@@ -226,7 +227,7 @@ pub(crate) fn recip<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Dot product accumulated layout
-pub(crate) fn dot<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn dot<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -337,7 +338,7 @@ pub(crate) fn dot<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Einsum
-pub(crate) fn einsum<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn einsum<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    inputs: &[ValTensor<F>],
@@ -524,14 +525,7 @@ pub(crate) fn einsum<F: PrimeField + TensorType + PartialOrd>(

                // Compute the product of all input tensors
                for pair in input_pairs {
-                    let product_across_pair = prod(
-                        config,
-                        region,
-                        &[pair.try_into().map_err(|e| {
-                            error!("{}", e);
-                            halo2_proofs::plonk::Error::Synthesis
-                        })?],
-                    )?;
+                    let product_across_pair = prod(config, region, &[pair.into()])?;

                    if let Some(product) = prod_res {
                        prod_res = Some(
@@ -563,7 +557,7 @@ pub(crate) fn einsum<F: PrimeField + TensorType + PartialOrd>(
    Ok(output)
 }

-fn _sort_ascending<F: PrimeField + TensorType + PartialOrd>(
+fn _sort_ascending<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -574,12 +568,12 @@ fn _sort_ascending<F: PrimeField + TensorType + PartialOrd>(
    let is_assigned = !input.any_unknowns()?;

    let sorted = if is_assigned {
-        input
-            .get_int_evals()?
-            .iter()
-            .sorted_by(|a, b| a.cmp(b))
-            .map(|x| Ok(Value::known(i128_to_felt(*x))))
-            .collect::<Result<Tensor<Value<F>>, Box<dyn Error>>>()?
+        let mut int_evals = input.get_int_evals()?;
+        int_evals.par_sort_unstable_by(|a, b| a.cmp(b));
+        int_evals
+            .par_iter()
+            .map(|x| Value::known(i128_to_felt(*x)))
+            .collect::<Tensor<Value<F>>>()
    } else {
        Tensor::new(
            Some(&vec![Value::<F>::unknown(); input.len()]),
@@ -607,7 +601,7 @@ fn _sort_ascending<F: PrimeField + TensorType + PartialOrd>(
 }

 ///
-fn _select_topk<F: PrimeField + TensorType + PartialOrd>(
+fn _select_topk<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -622,7 +616,7 @@ fn _select_topk<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Select top k elements
-pub(crate) fn topk_axes<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn topk_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -642,11 +636,12 @@ pub(crate) fn topk_axes<F: PrimeField + TensorType + PartialOrd>(
    Ok(output)
 }

-fn select<F: PrimeField + TensorType + PartialOrd>(
+fn select<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
 ) -> Result<ValTensor<F>, Box<dyn Error>> {
+    let start = instant::Instant::now();
    let (mut input, index) = (values[0].clone(), values[1].clone());
    input.flatten();

@@ -656,12 +651,13 @@ fn select<F: PrimeField + TensorType + PartialOrd>(

    let is_assigned = !input.any_unknowns()? && !index.any_unknowns()?;

-    let output: ValTensor<F> = if is_assigned {
+    let output: ValTensor<F> = if is_assigned && region.witness_gen() {
+        let felt_evals = input.get_felt_evals()?;
        index
            .get_int_evals()?
-            .iter()
-            .map(|x| Ok(Value::known(input.get_felt_evals()?.get(&[*x as usize]))))
-            .collect::<Result<Tensor<Value<F>>, Box<dyn Error>>>()?
+            .par_iter()
+            .map(|x| Value::known(felt_evals.get(&[*x as usize])))
+            .collect::<Tensor<Value<F>>>()
    } else {
        Tensor::new(
            Some(&vec![Value::<F>::unknown(); index.len()]),
@@ -673,10 +669,13 @@ fn select<F: PrimeField + TensorType + PartialOrd>(
    let (_, assigned_output) =
        dynamic_lookup(config, region, &[index, output], &[dim_indices, input])?;

+    let end = start.elapsed();
+    trace!("select took: {:?}", end);
+
    Ok(assigned_output)
 }

-fn one_hot<F: PrimeField + TensorType + PartialOrd>(
+fn one_hot<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -692,7 +691,7 @@ fn one_hot<F: PrimeField + TensorType + PartialOrd>(
    let output: ValTensor<F> = if is_assigned {
        let int_evals = input.get_int_evals()?;
        let res = tensor::ops::one_hot(&int_evals, num_classes, 1)?;
-        res.iter()
+        res.par_iter()
            .map(|x| Value::known(i128_to_felt(*x)))
            .collect::<Tensor<_>>()
    } else {
@@ -728,12 +727,13 @@ fn one_hot<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Dynamic lookup
-pub(crate) fn dynamic_lookup<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn dynamic_lookup<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    lookups: &[ValTensor<F>; 2],
    tables: &[ValTensor<F>; 2],
 ) -> Result<(ValTensor<F>, ValTensor<F>), Box<dyn Error>> {
+    let start = instant::Instant::now();
    // if not all lookups same length err
    if lookups[0].len() != lookups[1].len() {
        return Err("lookups must be same length".into());
@@ -753,20 +753,28 @@ pub(crate) fn dynamic_lookup<F: PrimeField + TensorType + PartialOrd>(
    let _table_1 = region.assign_dynamic_lookup(&config.dynamic_lookups.tables[1], &table_1)?;
    let table_len = table_0.len();

+    trace!("assigning tables took: {:?}", start.elapsed());
+
    // now create a vartensor of constants for the dynamic lookup index
    let table_index = create_constant_tensor(F::from(dynamic_lookup_index as u64), table_len);
    let _table_index =
        region.assign_dynamic_lookup(&config.dynamic_lookups.tables[2], &table_index)?;

+    trace!("assigning table index took: {:?}", start.elapsed());
+
    let lookup_0 = region.assign(&config.dynamic_lookups.inputs[0], &lookup_0)?;
    let lookup_1 = region.assign(&config.dynamic_lookups.inputs[1], &lookup_1)?;
    let lookup_len = lookup_0.len();

+    trace!("assigning lookups took: {:?}", start.elapsed());
+
    // now set the lookup index
    let lookup_index = create_constant_tensor(F::from(dynamic_lookup_index as u64), lookup_len);

    let _lookup_index = region.assign(&config.dynamic_lookups.inputs[2], &lookup_index)?;

+    trace!("assigning lookup index took: {:?}", start.elapsed());
+
    if !region.is_dummy() {
        (0..table_len)
            .map(|i| {
@@ -802,11 +810,14 @@ pub(crate) fn dynamic_lookup<F: PrimeField + TensorType + PartialOrd>(
    region.increment_dynamic_lookup_index(1);
    region.increment(lookup_len);

+    let end = start.elapsed();
+    trace!("dynamic lookup took: {:?}", end);
+
    Ok((lookup_0, lookup_1))
 }

 /// Shuffle arg
-pub(crate) fn shuffles<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn shuffles<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    input: &[ValTensor<F>; 1],
@@ -869,7 +880,7 @@ pub(crate) fn shuffles<F: PrimeField + TensorType + PartialOrd>(
 }

 /// One hot accumulated layout
-pub(crate) fn one_hot_axis<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn one_hot_axis<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -922,7 +933,7 @@ pub(crate) fn one_hot_axis<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Gather accumulated layout
-pub(crate) fn gather<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn gather<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -950,7 +961,7 @@ pub(crate) fn gather<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Gather accumulated layout
-pub(crate) fn gather_elements<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn gather_elements<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -973,7 +984,7 @@ pub(crate) fn gather_elements<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Gather accumulated layout
-pub(crate) fn gather_nd<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn gather_nd<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -1024,7 +1035,7 @@ pub(crate) fn gather_nd<F: PrimeField + TensorType + PartialOrd>(
 /// Takes a tensor representing a multi-dimensional index and returns a tensor representing the linearized index.
 /// The linearized index is the index of the element in the flattened tensor.
 /// FOr instance if the dims is [3,5,2], the linearized index of [2] at dim 1 is 2*5 + 3 = 13
-pub(crate) fn linearize_element_index<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn linearize_element_index<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1032,6 +1043,7 @@ pub(crate) fn linearize_element_index<F: PrimeField + TensorType + PartialOrd>(
    dim: usize,
    is_flat_index: bool,
 ) -> Result<ValTensor<F>, Box<dyn Error>> {
+    let start_time = instant::Instant::now();
    let index = values[0].clone();
    if !is_flat_index {
        assert_eq!(index.dims().len(), dims.len());
@@ -1105,6 +1117,9 @@ pub(crate) fn linearize_element_index<F: PrimeField + TensorType + PartialOrd>(

    region.apply_in_loop(&mut output, inner_loop_function)?;

+    let elapsed = start_time.elapsed();
+    trace!("linearize_element_index took: {:?}", elapsed);
+
    Ok(output.into())
 }

@@ -1125,7 +1140,7 @@ pub(crate) fn linearize_element_index<F: PrimeField + TensorType + PartialOrd>(
 ///     If indices_shape[-1] == r-b, since the rank of indices is q, indices can be thought of as N (q-b-1)-dimensional tensors containing 1-D tensors of dimension r-b, where N is an integer equals to the product of 1 and all the elements in the batch dimensions of the indices_shape.
 ///     Let us think of each such r-b ranked tensor as indices_slice. Each scalar value corresponding to data[0:b-1,indices_slice] is filled into the corresponding location of the (q-b-1)-dimensional tensor to form the output tensor (Example 1 below)
 ///     If indices_shape[-1] < r-b, since the rank of indices is q, indices can be thought of as N (q-b-1)-dimensional tensor containing 1-D tensors of dimension < r-b. Let us think of each such tensors as indices_slice. Each tensor slice corresponding to data[0:b-1, indices_slice , :] is filled into the corresponding location of the (q-b-1)-dimensional tensor to form the output tensor (Examples 2, 3, 4 and 5 below)
-pub(crate) fn linearize_nd_index<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn linearize_nd_index<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1232,7 +1247,6 @@ pub(crate) fn linearize_nd_index<F: PrimeField + TensorType + PartialOrd>(
            const_offset += F::from(coord[i] as u64) * dim_multiplier[i];
        }

-
        let const_offset = create_constant_tensor(const_offset, 1);

        let mut results = vec![];
@@ -1250,16 +1264,18 @@ pub(crate) fn linearize_nd_index<F: PrimeField + TensorType + PartialOrd>(
            let res = sum(config, region, &[res])?;
            results.push(res.get_inner_tensor()?.clone());
            // assert than res is less than the product of the dims
-            assert!(
+            if region.witness_gen() {
+                assert!(
                res.get_int_evals()?
                    .iter()
                    .all(|x| *x < dims.iter().product::<usize>() as i128),
                "res is greater than the product of the dims {} (coord={}, index_dim_multiplier={}, res={})",
                dims.iter().product::<usize>(),
-                index_val.show(), 
+                index_val.show(),
                index_dim_multiplier.show(),
                res.show()
            );
+            }
        }

        let result_tensor = Tensor::from(results.into_iter());
@@ -1273,7 +1289,9 @@ pub(crate) fn linearize_nd_index<F: PrimeField + TensorType + PartialOrd>(
    Ok(output.into())
 }

-pub(crate) fn get_missing_set_elements<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn get_missing_set_elements<
+    F: PrimeField + TensorType + PartialOrd + std::hash::Hash,
+>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -1304,7 +1322,7 @@ pub(crate) fn get_missing_set_elements<F: PrimeField + TensorType + PartialOrd>(
        }

        fullset_evals
-            .iter()
+            .par_iter()
            .map(|x| Value::known(i128_to_felt(*x)))
            .collect::<Tensor<Value<F>>>()
            .into()
@@ -1337,7 +1355,7 @@ pub(crate) fn get_missing_set_elements<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Gather accumulated layout
-pub(crate) fn scatter_elements<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn scatter_elements<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 3],
@@ -1354,14 +1372,14 @@ pub(crate) fn scatter_elements<F: PrimeField + TensorType + PartialOrd>(

    let is_assigned = !input.any_unknowns()? && !index.any_unknowns()? && !src.any_unknowns()?;

-    let claimed_output: ValTensor<F> = if is_assigned {
+    let claimed_output: ValTensor<F> = if is_assigned && region.witness_gen() {
        let input_inner = input.get_int_evals()?;
        let index_inner = index.get_int_evals()?.map(|x| x as usize);
        let src_inner = src.get_int_evals()?;

        let res = tensor::ops::scatter(&input_inner, &index_inner, &src_inner, dim)?;

-        res.iter()
+        res.par_iter()
            .map(|x| Value::known(i128_to_felt(*x)))
            .collect::<Tensor<Value<F>>>()
            .into()
@@ -1419,7 +1437,7 @@ pub(crate) fn scatter_elements<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Scatter Nd
-pub(crate) fn scatter_nd<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn scatter_nd<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 3],
@@ -1433,14 +1451,14 @@ pub(crate) fn scatter_nd<F: PrimeField + TensorType + PartialOrd>(

    let is_assigned = !input.any_unknowns()? && !index.any_unknowns()? && !src.any_unknowns()?;

-    let claimed_output: ValTensor<F> = if is_assigned {
+    let claimed_output: ValTensor<F> = if is_assigned && region.witness_gen() {
        let input_inner = input.get_int_evals()?;
        let index_inner = index.get_int_evals()?.map(|x| x as usize);
        let src_inner = src.get_int_evals()?;

        let res = tensor::ops::scatter_nd(&input_inner, &index_inner, &src_inner)?;

-        res.iter()
+        res.par_iter()
            .map(|x| Value::known(i128_to_felt(*x)))
            .collect::<Tensor<Value<F>>>()
            .into()
@@ -1457,7 +1475,6 @@ pub(crate) fn scatter_nd<F: PrimeField + TensorType + PartialOrd>(
    region.increment(claimed_output.len());
    claimed_output.reshape(input.dims())?;

-
    // scatter elements is the inverse of gather elements
    let (gather_src, linear_index) =
        gather_nd(config, region, &[claimed_output.clone(), index.clone()], 0)?;
@@ -1498,7 +1515,7 @@ pub(crate) fn scatter_nd<F: PrimeField + TensorType + PartialOrd>(
 }

 /// sum accumulated layout
-pub(crate) fn sum<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn sum<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1581,7 +1598,7 @@ pub(crate) fn sum<F: PrimeField + TensorType + PartialOrd>(
 }

 /// product accumulated layout
-pub(crate) fn prod<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn prod<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1661,7 +1678,7 @@ pub(crate) fn prod<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Axes wise op wrapper
-fn axes_wise_op<F: PrimeField + TensorType + PartialOrd>(
+fn axes_wise_op<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1722,7 +1739,7 @@ fn axes_wise_op<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Sum accumulated layout
-pub(crate) fn prod_axes<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn prod_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1733,7 +1750,7 @@ pub(crate) fn prod_axes<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Sum accumulated layout
-pub(crate) fn sum_axes<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn sum_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1744,7 +1761,7 @@ pub(crate) fn sum_axes<F: PrimeField + TensorType + PartialOrd>(
 }

 /// argmax layout
-pub(crate) fn argmax_axes<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn argmax_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1762,7 +1779,7 @@ pub(crate) fn argmax_axes<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Max accumulated layout
-pub(crate) fn max_axes<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn max_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1774,7 +1791,7 @@ pub(crate) fn max_axes<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Argmin layout
-pub(crate) fn argmin_axes<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn argmin_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1792,7 +1809,7 @@ pub(crate) fn argmin_axes<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Min accumulated layout
-pub(crate) fn min_axes<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn min_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1804,7 +1821,7 @@ pub(crate) fn min_axes<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Pairwise (elementwise) op layout
-pub(crate) fn pairwise<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn pairwise<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -1958,8 +1975,23 @@ pub(crate) fn pairwise<F: PrimeField + TensorType + PartialOrd>(
    Ok(output)
 }

+pub(crate) fn mean_of_squares_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
+    config: &BaseConfig<F>,
+    region: &mut RegionCtx<F>,
+    values: &[ValTensor<F>; 1],
+    axes: &[usize],
+) -> Result<ValTensor<F>, Box<dyn Error>> {
+    let squared = pow(config, region, values, 2)?;
+    let sum_squared = sum_axes(config, region, &[squared], axes)?;
+
+    let dividand: usize = values[0].len() / sum_squared.len();
+
+    let mean_squared = div(config, region, &[sum_squared], F::from(dividand as u64))?;
+    Ok(mean_squared)
+}
+
 /// expand the tensor to the given shape
-pub(crate) fn expand<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn expand<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -1972,7 +2004,7 @@ pub(crate) fn expand<F: PrimeField + TensorType + PartialOrd>(
 }

 ///
-pub(crate) fn greater<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn greater<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -1995,7 +2027,7 @@ pub(crate) fn greater<F: PrimeField + TensorType + PartialOrd>(
 }

 ///
-pub(crate) fn greater_equal<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn greater_equal<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -2018,7 +2050,7 @@ pub(crate) fn greater_equal<F: PrimeField + TensorType + PartialOrd>(
 }

 ///
-pub(crate) fn less<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn less<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -2028,7 +2060,7 @@ pub(crate) fn less<F: PrimeField + TensorType + PartialOrd>(
 }

 ///
-pub(crate) fn less_equal<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn less_equal<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -2038,7 +2070,7 @@ pub(crate) fn less_equal<F: PrimeField + TensorType + PartialOrd>(
 }

 /// And boolean operation
-pub(crate) fn and<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn and<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -2049,7 +2081,7 @@ pub(crate) fn and<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Or boolean operation
-pub(crate) fn or<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn or<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -2065,7 +2097,7 @@ pub(crate) fn or<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Equality boolean operation
-pub(crate) fn equals<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn equals<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -2075,7 +2107,7 @@ pub(crate) fn equals<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Equality boolean operation
-pub(crate) fn equals_zero<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn equals_zero<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2109,7 +2141,7 @@ pub(crate) fn equals_zero<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Xor boolean operation
-pub(crate) fn xor<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn xor<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -2135,7 +2167,7 @@ pub(crate) fn xor<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Not boolean operation
-pub(crate) fn not<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn not<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2151,7 +2183,7 @@ pub(crate) fn not<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Iff
-pub(crate) fn iff<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn iff<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 3],
@@ -2175,7 +2207,7 @@ pub(crate) fn iff<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Negation operation accumulated layout
-pub(crate) fn neg<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn neg<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2185,7 +2217,7 @@ pub(crate) fn neg<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Sumpool accumulated layout
-pub(crate) fn sumpool<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn sumpool<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>],
@@ -2239,7 +2271,7 @@ pub(crate) fn sumpool<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Convolution accumulated layout
-pub(crate) fn max_pool2d<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn max_pool2d<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2304,7 +2336,7 @@ pub(crate) fn max_pool2d<F: PrimeField + TensorType + PartialOrd>(

 /// DeConvolution accumulated layout
 pub(crate) fn deconv<
-    F: PrimeField + TensorType + PartialOrd + std::marker::Send + std::marker::Sync,
+    F: PrimeField + TensorType + PartialOrd + std::hash::Hash + std::marker::Send + std::marker::Sync,
 >(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
@@ -2397,7 +2429,7 @@ pub(crate) fn deconv<

 /// Convolution accumulated layout
 pub(crate) fn conv<
-    F: PrimeField + TensorType + PartialOrd + std::marker::Send + std::marker::Sync,
+    F: PrimeField + TensorType + PartialOrd + std::hash::Hash + std::marker::Send + std::marker::Sync,
 >(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
@@ -2578,7 +2610,7 @@ pub(crate) fn conv<
 }

 /// Power accumulated layout
-pub(crate) fn pow<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn pow<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2594,7 +2626,7 @@ pub(crate) fn pow<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Rescaled op accumulated layout
-pub(crate) fn rescale<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn rescale<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>],
@@ -2616,7 +2648,7 @@ pub(crate) fn rescale<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Dummy (no contraints) reshape layout
-pub(crate) fn reshape<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn reshape<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    values: &[ValTensor<F>; 1],
    new_dims: &[usize],
 ) -> Result<ValTensor<F>, Box<dyn Error>> {
@@ -2626,7 +2658,7 @@ pub(crate) fn reshape<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Dummy (no contraints) move_axis layout
-pub(crate) fn move_axis<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn move_axis<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    values: &[ValTensor<F>; 1],
    source: usize,
    destination: usize,
@@ -2637,7 +2669,7 @@ pub(crate) fn move_axis<F: PrimeField + TensorType + PartialOrd>(
 }

 /// resize layout
-pub(crate) fn resize<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn resize<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2651,7 +2683,7 @@ pub(crate) fn resize<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Slice layout
-pub(crate) fn slice<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn slice<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2674,7 +2706,7 @@ pub(crate) fn slice<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Trilu layout
-pub(crate) fn trilu<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn trilu<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2698,7 +2730,7 @@ pub(crate) fn trilu<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Concat layout
-pub(crate) fn concat<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn concat<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    values: &[ValTensor<F>],
    axis: &usize,
 ) -> Result<ValTensor<F>, Box<dyn Error>> {
@@ -2710,7 +2742,7 @@ pub(crate) fn concat<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Identity constraint. Usually used to constrain an instance column to an advice so the returned cells / values can be operated upon.
-pub(crate) fn identity<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn identity<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2725,7 +2757,7 @@ pub(crate) fn identity<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Boolean identity constraint. Usually used to constrain an instance column to an advice so the returned cells / values can be operated upon.
-pub(crate) fn boolean_identity<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn boolean_identity<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2761,7 +2793,7 @@ pub(crate) fn boolean_identity<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Downsample layout
-pub(crate) fn downsample<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn downsample<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2778,7 +2810,7 @@ pub(crate) fn downsample<F: PrimeField + TensorType + PartialOrd>(
 }

 /// layout for enforcing two sets of cells to be equal
-pub(crate) fn enforce_equality<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn enforce_equality<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
@@ -2804,7 +2836,7 @@ pub(crate) fn enforce_equality<F: PrimeField + TensorType + PartialOrd>(
 }

 /// layout for range check.
-pub(crate) fn range_check<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn range_check<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2860,12 +2892,13 @@ pub(crate) fn range_check<F: PrimeField + TensorType + PartialOrd>(
            .collect::<Result<Vec<_>, Box<dyn Error>>>()?;
    }

-    if region.throw_range_check_error() {
+    let is_assigned = !w.any_unknowns()?;
+    if is_assigned && region.witness_gen() {
        // assert is within range
        let int_values = w.get_int_evals()?;
-        for v in int_values {
-            if v < range.0 || v > range.1 {
-                log::debug!("Value ({:?}) out of range: {:?}", v, range);
+        for v in int_values.iter() {
+            if v < &range.0 || v > &range.1 {
+                log::error!("Value ({:?}) out of range: {:?}", v, range);
                return Err(Box::new(TensorError::TableLookupError));
            }
        }
@@ -2885,7 +2918,7 @@ pub(crate) fn range_check<F: PrimeField + TensorType + PartialOrd>(
 }

 /// layout for nonlinearity check.
-pub(crate) fn nonlinearity<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn nonlinearity<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -2987,7 +3020,7 @@ pub(crate) fn nonlinearity<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Argmax
-pub(crate) fn argmax<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn argmax<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -3023,7 +3056,7 @@ pub(crate) fn argmax<F: PrimeField + TensorType + PartialOrd>(
 }

 /// Argmin
-pub(crate) fn argmin<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn argmin<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -3059,7 +3092,7 @@ pub(crate) fn argmin<F: PrimeField + TensorType + PartialOrd>(
 }

 /// max layout
-pub(crate) fn max<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn max<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -3069,7 +3102,7 @@ pub(crate) fn max<F: PrimeField + TensorType + PartialOrd>(
 }

 /// min layout
-pub(crate) fn min<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn min<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -3077,7 +3110,7 @@ pub(crate) fn min<F: PrimeField + TensorType + PartialOrd>(
    _sort_ascending(config, region, values)?.get_slice(&[0..1])
 }

-fn multi_dim_axes_op<F: PrimeField + TensorType + PartialOrd>(
+fn multi_dim_axes_op<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
@@ -3180,18 +3213,19 @@ fn multi_dim_axes_op<F: PrimeField + TensorType + PartialOrd>(
 }

 /// softmax layout
-pub(crate) fn softmax_axes<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn softmax_axes<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
-    scale: utils::F32,
+    input_scale: utils::F32,
+    output_scale: utils::F32,
    axes: &[usize],
 ) -> Result<ValTensor<F>, Box<dyn Error>> {
    let soft_max_at_scale = move |config: &BaseConfig<F>,
                                  region: &mut RegionCtx<F>,
                                  values: &[ValTensor<F>; 1]|
          -> Result<ValTensor<F>, Box<dyn Error>> {
-        softmax(config, region, values, scale)
+        softmax(config, region, values, input_scale, output_scale)
    };

    let output = multi_dim_axes_op(config, region, values, axes, soft_max_at_scale)?;
@@ -3199,33 +3233,66 @@ pub(crate) fn softmax_axes<F: PrimeField + TensorType + PartialOrd>(
    Ok(output)
 }

-/// softmax func
-pub(crate) fn softmax<F: PrimeField + TensorType + PartialOrd>(
+/// percent func
+pub(crate) fn percent<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 1],
-    scale: utils::F32,
+    input_scale: utils::F32,
+    output_scale: utils::F32,
 ) -> Result<ValTensor<F>, Box<dyn Error>> {
-    // elementwise exponential
-    let ex = nonlinearity(config, region, values, &LookupOp::Exp { scale })?;
-
+    let is_assigned = values[0].all_prev_assigned();
+    let mut input = values[0].clone();
+    if !is_assigned {
+        input = region.assign(&config.custom_gates.inputs[0], &values[0])?;
+        region.increment(input.len());
+    };
    // sum of exps
-    let denom = sum(config, region, &[ex.clone()])?;
-    // get the inverse
-
-    let felt_scale = F::from(scale.0 as u64);
-    let inv_denom = recip(config, region, &[denom], felt_scale, felt_scale)?;
+    let denom = sum(config, region, &[input.clone()])?;

+    let input_felt_scale = F::from(input_scale.0 as u64);
+    let output_felt_scale = F::from(output_scale.0 as u64);
+    let inv_denom = recip(
+        config,
+        region,
+        &[denom],
+        input_felt_scale,
+        output_felt_scale,
+    )?;
    // product of num * (1 / denom) = 2*output_scale
-    let softmax = pairwise(config, region, &[ex, inv_denom], BaseOp::Mult)?;
+    let percent = pairwise(config, region, &[input, inv_denom], BaseOp::Mult)?;

-    Ok(softmax)
+    // rebase the percent to 2x the scale
+    loop_div(config, region, &[percent], input_felt_scale)
+}
+
+/// softmax func
+pub(crate) fn softmax<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
+    config: &BaseConfig<F>,
+    region: &mut RegionCtx<F>,
+    values: &[ValTensor<F>; 1],
+    input_scale: utils::F32,
+    output_scale: utils::F32,
+) -> Result<ValTensor<F>, Box<dyn Error>> {
+    // get the max then subtract it
+    let max_val = max(config, region, values)?;
+    // rebase the input to 0
+    let sub = pairwise(config, region, &[values[0].clone(), max_val], BaseOp::Sub)?;
+    // elementwise exponential
+    let ex = nonlinearity(
+        config,
+        region,
+        &[sub],
+        &LookupOp::Exp { scale: input_scale },
+    )?;
+
+    percent(config, region, &[ex.clone()], input_scale, output_scale)
 }

 /// Checks that the percent error between the expected public output and the actual output value
 /// is within the percent error expressed by the `tol` input, where `tol == 1.0` means the percent
 /// error tolerance is 1 percent.
-pub(crate) fn range_check_percent<F: PrimeField + TensorType + PartialOrd>(
+pub(crate) fn range_check_percent<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    config: &BaseConfig<F>,
    region: &mut RegionCtx<F>,
    values: &[ValTensor<F>; 2],
--- a/src/circuit/ops/lookup.rs
+++ b/src/circuit/ops/lookup.rs
@@ -137,7 +137,7 @@ impl LookupOp {
    }
 }

-impl<F: PrimeField + TensorType + PartialOrd> Op<F> for LookupOp {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for LookupOp {
    /// Returns a reference to the Any trait.
    fn as_any(&self) -> &dyn Any {
        self
--- a/src/circuit/ops/mod.rs
+++ b/src/circuit/ops/mod.rs
@@ -27,12 +27,14 @@ pub mod region;

 /// A struct representing the result of a forward pass.
 #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
-pub struct ForwardResult<F: PrimeField + TensorType + PartialOrd> {
+pub struct ForwardResult<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> {
    pub(crate) output: Tensor<F>,
 }

 /// A trait representing operations that can be represented as constraints in a circuit.
-pub trait Op<F: PrimeField + TensorType + PartialOrd>: std::fmt::Debug + Send + Sync + Any {
+pub trait Op<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>:
+    std::fmt::Debug + Send + Sync + Any
+{
    /// Matches a [Op] to an operation in the `tensor::ops` module.
    fn f(&self, x: &[Tensor<F>]) -> Result<ForwardResult<F>, TensorError>;
    /// Returns a string representation of the operation.
@@ -98,7 +100,7 @@ pub trait Op<F: PrimeField + TensorType + PartialOrd>: std::fmt::Debug + Send +
    }
 }

-impl<F: PrimeField + TensorType + PartialOrd> Clone for Box<dyn Op<F>> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Clone for Box<dyn Op<F>> {
    fn clone(&self) -> Self {
        self.clone_dyn()
    }
@@ -165,7 +167,7 @@ pub struct Input {
    pub datum_type: InputType,
 }

-impl<F: PrimeField + TensorType + PartialOrd> Op<F> for Input {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Input {
    fn out_scale(&self, _: Vec<crate::Scale>) -> Result<crate::Scale, Box<dyn Error>> {
        Ok(self.scale)
    }
@@ -226,7 +228,7 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for Input {
 #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct Unknown;

-impl<F: PrimeField + TensorType + PartialOrd> Op<F> for Unknown {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Op<F> for Unknown {
    fn out_scale(&self, _: Vec<crate::Scale>) -> Result<crate::Scale, Box<dyn Error>> {
        Ok(0)
    }
@@ -256,7 +258,7 @@ impl<F: PrimeField + TensorType + PartialOrd> Op<F> for Unknown {

 ///
 #[derive(Clone, Debug, Serialize, Deserialize)]
-pub struct Constant<F: PrimeField + TensorType + PartialOrd> {
+pub struct Constant<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> {
    ///
    pub quantized_values: Tensor<F>,
    ///
@@ -266,7 +268,7 @@ pub struct Constant<F: PrimeField + TensorType + PartialOrd> {
    pub pre_assigned_val: Option<ValTensor<F>>,
 }

-impl<F: PrimeField + TensorType + PartialOrd> Constant<F> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Constant<F> {
    ///
    pub fn new(quantized_values: Tensor<F>, raw_values: Tensor<f32>) -> Self {
        Self {
@@ -293,8 +295,14 @@ impl<F: PrimeField + TensorType + PartialOrd> Constant<F> {
    }
 }

-impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<'de>> Op<F>
-    for Constant<F>
+impl<
+        F: PrimeField
+            + TensorType
+            + PartialOrd
+            + std::hash::Hash
+            + Serialize
+            + for<'de> Deserialize<'de>,
+    > Op<F> for Constant<F>
 {
    fn as_any(&self) -> &dyn Any {
        self
--- a/src/circuit/ops/poly.rs
+++ b/src/circuit/ops/poly.rs
@@ -1,6 +1,6 @@
 use crate::{
    circuit::layouts,
-    fieldutils::felt_to_i128,
+    fieldutils::{felt_to_i128, i128_to_felt},
    tensor::{self, Tensor, TensorError},
 };

@@ -62,6 +62,9 @@ pub enum PolyOp {
    Sum {
        axes: Vec<usize>,
    },
+    MeanOfSquares {
+        axes: Vec<usize>,
+    },
    Prod {
        axes: Vec<usize>,
        len_prod: usize,
@@ -89,8 +92,14 @@ pub enum PolyOp {
    },
 }

-impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<'de>> Op<F>
-    for PolyOp
+impl<
+        F: PrimeField
+            + TensorType
+            + PartialOrd
+            + std::hash::Hash
+            + Serialize
+            + for<'de> Deserialize<'de>,
+    > Op<F> for PolyOp
 {
    /// Returns a reference to the Any trait.
    fn as_any(&self) -> &dyn Any {
@@ -99,10 +108,28 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<

    fn as_string(&self) -> String {
        match &self {
-            PolyOp::GatherElements { dim, .. } => format!("GATHERELEMENTS (dim={})", dim),
-            PolyOp::GatherND { batch_dims, .. } => format!("GATHERND (batch_dims={})", batch_dims),
-            PolyOp::ScatterElements { dim, .. } => format!("SCATTERELEMENTS (dim={})", dim),
-            PolyOp::ScatterND { .. } => "SCATTERND".into(),
+            PolyOp::GatherElements { dim, constant_idx } => format!(
+                "GATHERELEMENTS (dim={}, constant_idx{})",
+                dim,
+                constant_idx.is_some()
+            ),
+            PolyOp::GatherND {
+                batch_dims,
+                indices,
+            } => format!(
+                "GATHERND (batch_dims={}, constant_idx{})",
+                batch_dims,
+                indices.is_some()
+            ),
+            PolyOp::MeanOfSquares { axes } => format!("MEANOFSQUARES (axes={:?})", axes),
+            PolyOp::ScatterElements { dim, constant_idx } => format!(
+                "SCATTERELEMENTS (dim={}, constant_idx{})",
+                dim,
+                constant_idx.is_some()
+            ),
+            PolyOp::ScatterND { constant_idx } => {
+                format!("SCATTERND (constant_idx={})", constant_idx.is_some())
+            }
            PolyOp::MultiBroadcastTo { shape } => format!("MULTIBROADCASTTO (shape={:?})", shape),
            PolyOp::MoveAxis { .. } => "MOVEAXIS".into(),
            PolyOp::Downsample { .. } => "DOWNSAMPLE".into(),
@@ -140,6 +167,10 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<
    fn f(&self, inputs: &[Tensor<F>]) -> Result<ForwardResult<F>, TensorError> {
        let mut inputs = inputs.to_vec();
        let res = match &self {
+            PolyOp::MeanOfSquares { axes } => {
+                let x = inputs[0].map(|x| felt_to_i128(x));
+                Ok(tensor::ops::nonlinearities::mean_of_squares_axes(&x, axes).map(i128_to_felt))
+            }
            PolyOp::MultiBroadcastTo { shape } => {
                if 1 != inputs.len() {
                    return Err(TensorError::DimMismatch(
@@ -286,6 +317,9 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<
            PolyOp::MultiBroadcastTo { shape } => {
                layouts::expand(config, region, values[..].try_into()?, shape)?
            }
+            PolyOp::MeanOfSquares { axes } => {
+                layouts::mean_of_squares_axes(config, region, values[..].try_into()?, axes)?
+            }
            PolyOp::Xor => layouts::xor(config, region, values[..].try_into()?)?,
            PolyOp::Or => layouts::or(config, region, values[..].try_into()?)?,
            PolyOp::And => layouts::and(config, region, values[..].try_into()?)?,
@@ -398,6 +432,7 @@ impl<F: PrimeField + TensorType + PartialOrd + Serialize + for<'de> Deserialize<

    fn out_scale(&self, in_scales: Vec<crate::Scale>) -> Result<crate::Scale, Box<dyn Error>> {
        let scale = match self {
+            PolyOp::MeanOfSquares { .. } => 2 * in_scales[0],
            PolyOp::Xor | PolyOp::Or | PolyOp::And | PolyOp::Not => 0,
            PolyOp::Iff => in_scales[1],
            PolyOp::Einsum { .. } => {
--- a/src/circuit/ops/region.rs
+++ b/src/circuit/ops/region.rs
@@ -2,24 +2,28 @@ use crate::{
    circuit::table::Range,
    tensor::{Tensor, TensorError, TensorType, ValTensor, ValType, VarTensor},
 };
+#[cfg(not(target_arch = "wasm32"))]
+use colored::Colorize;
 use halo2_proofs::{
    circuit::Region,
    plonk::{Error, Selector},
 };
 use halo2curves::ff::PrimeField;
+use portable_atomic::AtomicI128 as AtomicInt;
 use std::{
    cell::RefCell,
-    collections::HashSet,
+    collections::{HashMap, HashSet},
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc, Mutex,
    },
 };

-use portable_atomic::AtomicI128 as AtomicInt;
-
 use super::lookup::LookupOp;

+/// Constants map
+pub type ConstantsMap<F> = HashMap<F, ValType<F>>;
+
 /// Dynamic lookup index
 #[derive(Clone, Debug, Default)]
 pub struct DynamicLookupIndex {
@@ -120,12 +124,11 @@ impl From<Box<dyn std::error::Error>> for RegionError {

 #[derive(Debug)]
 /// A context for a region
-pub struct RegionCtx<'a, F: PrimeField + TensorType + PartialOrd> {
+pub struct RegionCtx<'a, F: PrimeField + TensorType + PartialOrd + std::hash::Hash> {
    region: Option<RefCell<Region<'a, F>>>,
    row: usize,
    linear_coord: usize,
    num_inner_cols: usize,
-    total_constants: usize,
    dynamic_lookup_index: DynamicLookupIndex,
    shuffle_index: ShuffleIndex,
    used_lookups: HashSet<LookupOp>,
@@ -133,13 +136,34 @@ pub struct RegionCtx<'a, F: PrimeField + TensorType + PartialOrd> {
    max_lookup_inputs: i128,
    min_lookup_inputs: i128,
    max_range_size: i128,
-    throw_range_check_error: bool,
+    witness_gen: bool,
+    assigned_constants: ConstantsMap<F>,
 }

-impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
+impl<'a, F: PrimeField + TensorType + PartialOrd + std::hash::Hash> RegionCtx<'a, F> {
+    #[cfg(not(target_arch = "wasm32"))]
    ///
-    pub fn increment_total_constants(&mut self, n: usize) {
-        self.total_constants += n;
+    pub fn debug_report(&self) {
+        log::debug!(
+            "(rows={}, coord={}, constants={}, max_lookup_inputs={}, min_lookup_inputs={}, max_range_size={}, dynamic_lookup_col_coord={}, shuffle_col_coord={})",
+            self.row().to_string().blue(),
+            self.linear_coord().to_string().yellow(),
+            self.total_constants().to_string().red(),
+            self.max_lookup_inputs().to_string().green(),
+            self.min_lookup_inputs().to_string().green(),
+            self.max_range_size().to_string().green(),
+            self.dynamic_lookup_col_coord().to_string().green(),
+            self.shuffle_col_coord().to_string().green());
+    }
+
+    ///
+    pub fn assigned_constants(&self) -> &ConstantsMap<F> {
+        &self.assigned_constants
+    }
+
+    ///
+    pub fn update_constants(&mut self, constants: ConstantsMap<F>) {
+        self.assigned_constants.extend(constants);
    }

    ///
@@ -163,8 +187,8 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
    }

    ///
-    pub fn throw_range_check_error(&self) -> bool {
-        self.throw_range_check_error
+    pub fn witness_gen(&self) -> bool {
+        self.witness_gen
    }

    /// Create a new region context
@@ -177,7 +201,6 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            num_inner_cols,
            row,
            linear_coord,
-            total_constants: 0,
            dynamic_lookup_index: DynamicLookupIndex::default(),
            shuffle_index: ShuffleIndex::default(),
            used_lookups: HashSet::new(),
@@ -185,9 +208,22 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            max_lookup_inputs: 0,
            min_lookup_inputs: 0,
            max_range_size: 0,
-            throw_range_check_error: false,
+            witness_gen: true,
+            assigned_constants: HashMap::new(),
        }
    }
+
+    /// Create a new region context
+    pub fn new_with_constants(
+        region: Region<'a, F>,
+        row: usize,
+        num_inner_cols: usize,
+        constants: ConstantsMap<F>,
+    ) -> RegionCtx<'a, F> {
+        let mut new_self = Self::new(region, row, num_inner_cols);
+        new_self.assigned_constants = constants;
+        new_self
+    }
    /// Create a new region context from a wrapped region
    pub fn from_wrapped_region(
        region: Option<RefCell<Region<'a, F>>>,
@@ -202,7 +238,6 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            num_inner_cols,
            linear_coord,
            row,
-            total_constants: 0,
            dynamic_lookup_index,
            shuffle_index,
            used_lookups: HashSet::new(),
@@ -210,16 +245,13 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            max_lookup_inputs: 0,
            min_lookup_inputs: 0,
            max_range_size: 0,
-            throw_range_check_error: false,
+            witness_gen: false,
+            assigned_constants: HashMap::new(),
        }
    }

    /// Create a new region context
-    pub fn new_dummy(
-        row: usize,
-        num_inner_cols: usize,
-        throw_range_check_error: bool,
-    ) -> RegionCtx<'a, F> {
+    pub fn new_dummy(row: usize, num_inner_cols: usize, witness_gen: bool) -> RegionCtx<'a, F> {
        let region = None;
        let linear_coord = row * num_inner_cols;

@@ -228,7 +260,6 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            num_inner_cols,
            linear_coord,
            row,
-            total_constants: 0,
            dynamic_lookup_index: DynamicLookupIndex::default(),
            shuffle_index: ShuffleIndex::default(),
            used_lookups: HashSet::new(),
@@ -236,17 +267,17 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            max_lookup_inputs: 0,
            min_lookup_inputs: 0,
            max_range_size: 0,
-            throw_range_check_error,
+            witness_gen,
+            assigned_constants: HashMap::new(),
        }
    }

    /// Create a new region context
-    pub fn new_dummy_with_constants(
+    pub fn new_dummy_with_linear_coord(
        row: usize,
        linear_coord: usize,
-        total_constants: usize,
        num_inner_cols: usize,
-        throw_range_check_error: bool,
+        witness_gen: bool,
    ) -> RegionCtx<'a, F> {
        let region = None;
        RegionCtx {
@@ -254,7 +285,6 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            num_inner_cols,
            linear_coord,
            row,
-            total_constants,
            dynamic_lookup_index: DynamicLookupIndex::default(),
            shuffle_index: ShuffleIndex::default(),
            used_lookups: HashSet::new(),
@@ -262,7 +292,8 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            max_lookup_inputs: 0,
            min_lookup_inputs: 0,
            max_range_size: 0,
-            throw_range_check_error,
+            witness_gen,
+            assigned_constants: HashMap::new(),
        }
    }

@@ -312,29 +343,27 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
    ) -> Result<(), RegionError> {
        let row = AtomicUsize::new(self.row());
        let linear_coord = AtomicUsize::new(self.linear_coord());
-        let constants = AtomicUsize::new(self.total_constants());
        let max_lookup_inputs = AtomicInt::new(self.max_lookup_inputs());
        let min_lookup_inputs = AtomicInt::new(self.min_lookup_inputs());
        let lookups = Arc::new(Mutex::new(self.used_lookups.clone()));
        let range_checks = Arc::new(Mutex::new(self.used_range_checks.clone()));
        let dynamic_lookup_index = Arc::new(Mutex::new(self.dynamic_lookup_index.clone()));
        let shuffle_index = Arc::new(Mutex::new(self.shuffle_index.clone()));
+        let constants = Arc::new(Mutex::new(self.assigned_constants.clone()));

        *output = output
            .par_enum_map(|idx, _| {
                // we kick off the loop with the current offset
                let starting_offset = row.load(Ordering::SeqCst);
                let starting_linear_coord = linear_coord.load(Ordering::SeqCst);
-                let starting_constants = constants.load(Ordering::SeqCst);
                // get inner value of the locked lookups

                // we need to make sure that the region is not shared between threads
-                let mut local_reg = Self::new_dummy_with_constants(
+                let mut local_reg = Self::new_dummy_with_linear_coord(
                    starting_offset,
                    starting_linear_coord,
-                    starting_constants,
                    self.num_inner_cols,
-                    self.throw_range_check_error,
+                    self.witness_gen,
                );
                let res = inner_loop_function(idx, &mut local_reg);
                // we update the offset and constants
@@ -343,10 +372,6 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
                    local_reg.linear_coord() - starting_linear_coord,
                    Ordering::SeqCst,
                );
-                constants.fetch_add(
-                    local_reg.total_constants() - starting_constants,
-                    Ordering::SeqCst,
-                );

                max_lookup_inputs.fetch_max(local_reg.max_lookup_inputs(), Ordering::SeqCst);
                min_lookup_inputs.fetch_min(local_reg.min_lookup_inputs(), Ordering::SeqCst);
@@ -362,11 +387,13 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
                // update the shuffle index
                let mut shuffle_index = shuffle_index.lock().unwrap();
                shuffle_index.update(&local_reg.shuffle_index);
+                // update the constants
+                let mut constants = constants.lock().unwrap();
+                constants.extend(local_reg.assigned_constants);

                res
            })
            .map_err(|e| RegionError::from(format!("dummy_loop: {:?}", e)))?;
-        self.total_constants = constants.into_inner();
        self.linear_coord = linear_coord.into_inner();
        #[allow(trivial_numeric_casts)]
        {
@@ -410,6 +437,14 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
            .map_err(|e| {
                RegionError::from(format!("dummy_loop: failed to get shuffle index: {:?}", e))
            })?;
+        self.assigned_constants = Arc::try_unwrap(constants)
+            .map_err(|e| {
+                RegionError::from(format!("dummy_loop: failed to get constants: {:?}", e))
+            })?
+            .into_inner()
+            .map_err(|e| {
+                RegionError::from(format!("dummy_loop: failed to get constants: {:?}", e))
+            })?;

        Ok(())
    }
@@ -435,7 +470,7 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
        range: Range,
    ) -> Result<(), Box<dyn std::error::Error>> {
        if range.0 > range.1 {
-            return Err("update_max_min_lookup_range: invalid range".into());
+            return Err(format!("update_max_min_lookup_range: invalid range {:?}", range).into());
        }

        let range_size = (range.1 - range.0).abs();
@@ -477,7 +512,7 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {

    /// Get the total number of constants
    pub fn total_constants(&self) -> usize {
-        self.total_constants
+        self.assigned_constants.len()
    }

    /// Get the dynamic lookup index
@@ -525,26 +560,24 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
        self.max_range_size
    }

-    /// Assign a constant value
-    pub fn assign_constant(&mut self, var: &VarTensor, value: F) -> Result<ValType<F>, Error> {
-        self.total_constants += 1;
-        if let Some(region) = &self.region {
-            let cell = var.assign_constant(&mut region.borrow_mut(), self.linear_coord, value)?;
-            Ok(cell.into())
-        } else {
-            Ok(value.into())
-        }
-    }
    /// Assign a valtensor to a vartensor
    pub fn assign(
        &mut self,
        var: &VarTensor,
        values: &ValTensor<F>,
    ) -> Result<ValTensor<F>, Error> {
-        self.total_constants += values.num_constants();
        if let Some(region) = &self.region {
-            var.assign(&mut region.borrow_mut(), self.linear_coord, values)
+            var.assign(
+                &mut region.borrow_mut(),
+                self.linear_coord,
+                values,
+                &mut self.assigned_constants,
+            )
        } else {
+            if !values.is_instance() {
+                let values_map = values.create_constants_map_iterator();
+                self.assigned_constants.extend(values_map);
+            }
            Ok(values.clone())
        }
    }
@@ -560,14 +593,18 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
        var: &VarTensor,
        values: &ValTensor<F>,
    ) -> Result<ValTensor<F>, Error> {
-        self.total_constants += values.num_constants();
        if let Some(region) = &self.region {
            var.assign(
                &mut region.borrow_mut(),
                self.combined_dynamic_shuffle_coord(),
                values,
+                &mut self.assigned_constants,
            )
        } else {
+            if !values.is_instance() {
+                let values_map = values.create_constants_map_iterator();
+                self.assigned_constants.extend(values_map);
+            }
            Ok(values.clone())
        }
    }
@@ -594,13 +631,20 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
                self.linear_coord,
                values,
                ommissions,
+                &mut self.assigned_constants,
            )
        } else {
-            self.total_constants += values.num_constants();
            let inner_tensor = values.get_inner_tensor().unwrap();
+            let mut values_map = values.create_constants_map();
+
            for o in ommissions {
-                self.total_constants -= inner_tensor.get_flat_index(**o).is_constant() as usize;
+                if let ValType::Constant(value) = inner_tensor.get_flat_index(**o) {
+                    values_map.remove(&value);
+                }
            }
+
+            self.assigned_constants.extend(values_map);
+
            Ok(values.clone())
        }
    }
@@ -615,24 +659,24 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
    ) -> Result<(ValTensor<F>, usize), Error> {
        if let Some(region) = &self.region {
            // duplicates every nth element to adjust for column overflow
-            let (res, len, total_assigned_constants) = var.assign_with_duplication(
+            let (res, len) = var.assign_with_duplication(
                &mut region.borrow_mut(),
                self.row,
                self.linear_coord,
                values,
                check_mode,
                single_inner_col,
+                &mut self.assigned_constants,
            )?;
-            self.total_constants += total_assigned_constants;
            Ok((res, len))
        } else {
-            let (_, len, total_assigned_constants) = var.dummy_assign_with_duplication(
+            let (_, len) = var.dummy_assign_with_duplication(
                self.row,
                self.linear_coord,
                values,
                single_inner_col,
+                &mut self.assigned_constants,
            )?;
-            self.total_constants += total_assigned_constants;
            Ok((values.clone(), len))
        }
    }
@@ -699,9 +743,4 @@ impl<'a, F: PrimeField + TensorType + PartialOrd> RegionCtx<'a, F> {
        }
        Ok(())
    }
-
-    /// increment constants
-    pub fn increment_constants(&mut self, n: usize) {
-        self.total_constants += n
-    }
 }
--- a/src/circuit/table.rs
+++ b/src/circuit/table.rs
@@ -98,7 +98,7 @@ pub struct Table<F: PrimeField> {
    _marker: PhantomData<F>,
 }

-impl<F: PrimeField + TensorType + PartialOrd> Table<F> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Table<F> {
    /// get column index given input
    pub fn get_col_index(&self, input: F) -> F {
        //    range is split up into chunks of size col_size, find the chunk that input is in
@@ -138,7 +138,7 @@ pub fn num_cols_required(range_len: i128, col_size: usize) -> usize {
    (range_len / (col_size as i128)) as usize + 1
 }

-impl<F: PrimeField + TensorType + PartialOrd> Table<F> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> Table<F> {
    /// Configures the table.
    pub fn configure(
        cs: &mut ConstraintSystem<F>,
@@ -275,7 +275,7 @@ pub struct RangeCheck<F: PrimeField> {
    _marker: PhantomData<F>,
 }

-impl<F: PrimeField + TensorType + PartialOrd> RangeCheck<F> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> RangeCheck<F> {
    /// get first_element of column
    pub fn get_first_element(&self, chunk: usize) -> F {
        let chunk = chunk as i128;
@@ -303,7 +303,7 @@ impl<F: PrimeField + TensorType + PartialOrd> RangeCheck<F> {
    }
 }

-impl<F: PrimeField + TensorType + PartialOrd> RangeCheck<F> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> RangeCheck<F> {
    /// Configures the table.
    pub fn configure(cs: &mut ConstraintSystem<F>, range: Range, logrows: usize) -> RangeCheck<F> {
        log::debug!("range check range: {:?}", range);
--- a/src/circuit/tests.rs
+++ b/src/circuit/tests.rs
@@ -1911,6 +1911,8 @@ mod add_with_overflow {

 #[cfg(test)]
 mod add_with_overflow_and_poseidon {
+    use std::collections::HashMap;
+
    use halo2curves::bn256::Fr;

    use crate::circuit::modules::{
@@ -1969,8 +1971,10 @@ mod add_with_overflow_and_poseidon {
            let poseidon_chip: PoseidonChip<PoseidonSpec, WIDTH, RATE, WIDTH> =
                PoseidonChip::new(config.poseidon.clone());

-            let assigned_inputs_a = poseidon_chip.layout(&mut layouter, &self.inputs[0..1], 0)?;
-            let assigned_inputs_b = poseidon_chip.layout(&mut layouter, &self.inputs[1..2], 1)?;
+            let assigned_inputs_a =
+                poseidon_chip.layout(&mut layouter, &self.inputs[0..1], 0, &mut HashMap::new())?;
+            let assigned_inputs_b =
+                poseidon_chip.layout(&mut layouter, &self.inputs[1..2], 1, &mut HashMap::new())?;

            layouter.assign_region(|| "_new_module", |_| Ok(()))?;

--- a/src/commands.rs
+++ b/src/commands.rs
@@ -444,7 +444,7 @@ pub enum Commands {
        disable_selector_compression: bool,
        /// commitment used
        #[arg(long, default_value = DEFAULT_COMMITMENT)]
-        commitment: Commitments,
+        commitment: Option<Commitments>,
    },
    /// Aggregates proofs :)
    Aggregate {
@@ -479,7 +479,7 @@ pub enum Commands {
        split_proofs: bool,
        /// commitment used
        #[arg(long, default_value = DEFAULT_COMMITMENT)]
-        commitment: Commitments,
+        commitment: Option<Commitments>,
    },
    /// Compiles a circuit from onnx to a simplified graph (einsum + other ops) and parameters as sets of field elements
    CompileCircuit {
@@ -726,7 +726,7 @@ pub enum Commands {
        logrows: u32,
        /// commitment
        #[arg(long, default_value = DEFAULT_COMMITMENT)]
-        commitment: Commitments,
+        commitment: Option<Commitments>,
    },
    #[cfg(not(target_arch = "wasm32"))]
    /// Deploys an evm verifier that is generated by ezkl
--- a/src/execute.rs
+++ b/src/execute.rs
@@ -24,6 +24,8 @@ use crate::pfsys::{
 use crate::pfsys::{save_vk, srs::*};
 use crate::tensor::TensorError;
 use crate::{Commitments, RunArgs};
+#[cfg(not(target_arch = "wasm32"))]
+use colored::Colorize;
 #[cfg(unix)]
 use gag::Gag;
 use halo2_proofs::dev::VerifyFailure;
@@ -337,7 +339,7 @@ pub async fn run(command: Commands) -> Result<String, Box<dyn Error>> {
            logrows,
            split_proofs,
            disable_selector_compression,
-            commitment,
+            commitment.into(),
        ),
        Commands::Aggregate {
            proof_path,
@@ -358,7 +360,7 @@ pub async fn run(command: Commands) -> Result<String, Box<dyn Error>> {
            logrows,
            check_mode,
            split_proofs,
-            commitment,
+            commitment.into(),
        )
        .map(|e| serde_json::to_string(&e).unwrap()),
        Commands::Verify {
@@ -382,7 +384,7 @@ pub async fn run(command: Commands) -> Result<String, Box<dyn Error>> {
            srs_path,
            logrows,
            reduced_srs,
-            commitment,
+            commitment.into(),
        )
        .map(|e| serde_json::to_string(&e).unwrap()),
        #[cfg(not(target_arch = "wasm32"))]
@@ -538,7 +540,7 @@ fn check_srs_hash(
    let path = get_srs_path(logrows, srs_path, commitment);
    let hash = get_file_hash(&path)?;

-    let predefined_hash = match { crate::srs_sha::PUBLIC_SRS_SHA256_HASHES.get(&logrows) } {
+    let predefined_hash = match crate::srs_sha::PUBLIC_SRS_SHA256_HASHES.get(&logrows) {
        Some(h) => h,
        None => return Err(format!("SRS (k={}) hash not found in public set", logrows).into()),
    };
@@ -584,7 +586,7 @@ pub(crate) async fn get_srs_cmd(
    } else if let Some(settings_p) = settings_path {
        if settings_p.exists() {
            let settings = GraphSettings::load(&settings_p)?;
-            settings.run_args.commitment
+            settings.run_args.commitment.into()
        } else {
            return Err(err_string.into());
        }
@@ -664,27 +666,23 @@ pub(crate) async fn gen_witness(

    // if any of the settings have kzg visibility then we need to load the srs

+    let commitment: Commitments = settings.run_args.commitment.into();
+
    let start_time = Instant::now();
    let witness = if settings.module_requires_polycommit() {
-        if get_srs_path(
-            settings.run_args.logrows,
-            srs_path.clone(),
-            settings.run_args.commitment,
-        )
-        .exists()
-        {
-            match settings.run_args.commitment {
+        if get_srs_path(settings.run_args.logrows, srs_path.clone(), commitment).exists() {
+            match Commitments::from(settings.run_args.commitment) {
                Commitments::KZG => {
                    let srs: ParamsKZG<Bn256> = load_params_prover::<KZGCommitmentScheme<Bn256>>(
                        srs_path.clone(),
                        settings.run_args.logrows,
-                        settings.run_args.commitment,
+                        commitment,
                    )?;
                    circuit.forward::<KZGCommitmentScheme<_>>(
                        &mut input,
                        vk.as_ref(),
                        Some(&srs),
-                        false,
+                        true,
                    )?
                }
                Commitments::IPA => {
@@ -692,22 +690,22 @@ pub(crate) async fn gen_witness(
                        load_params_prover::<IPACommitmentScheme<G1Affine>>(
                            srs_path.clone(),
                            settings.run_args.logrows,
-                            settings.run_args.commitment,
+                            commitment,
                        )?;
                    circuit.forward::<IPACommitmentScheme<_>>(
                        &mut input,
                        vk.as_ref(),
                        Some(&srs),
-                        false,
+                        true,
                    )?
                }
            }
        } else {
            warn!("SRS for poly commit does not exist (will be ignored)");
-            circuit.forward::<KZGCommitmentScheme<Bn256>>(&mut input, vk.as_ref(), None, false)?
+            circuit.forward::<KZGCommitmentScheme<Bn256>>(&mut input, vk.as_ref(), None, true)?
        }
    } else {
-        circuit.forward::<KZGCommitmentScheme<Bn256>>(&mut input, vk.as_ref(), None, false)?
+        circuit.forward::<KZGCommitmentScheme<Bn256>>(&mut input, vk.as_ref(), None, true)?
    };

    // print each variable tuple (symbol, value) as symbol=value
@@ -819,7 +817,15 @@ impl AccuracyResults {
            let error = (original.clone() - calibrated.clone())?;
            let abs_error = error.map(|x| x.abs());
            let squared_error = error.map(|x| x.powi(2));
-            let percentage_error = error.enum_map(|i, x| Ok::<_, TensorError>(x / original[i]))?;
+            let percentage_error = error.enum_map(|i, x| {
+                // if everything is 0 then we can't divide by 0 so we just return 0
+                let res = if original[i] == 0.0 && x == 0.0 {
+                    0.0
+                } else {
+                    x / original[i]
+                };
+                Ok::<f32, TensorError>(res)
+            })?;
            let abs_percentage_error = percentage_error.map(|x| x.abs());

            errors.extend(error);
@@ -888,6 +894,7 @@ pub(crate) fn calibrate(
    only_range_check_rebase: bool,
    max_logrows: Option<u32>,
 ) -> Result<GraphSettings, Box<dyn Error>> {
+    use log::error;
    use std::collections::HashMap;
    use tabled::Table;

@@ -900,9 +907,9 @@ pub(crate) fn calibrate(
    let model = Model::from_run_args(&settings.run_args, &model_path)?;

    let chunks = data.split_into_batches(model.graph.input_shapes()?)?;
-    info!("num of calibration batches: {}", chunks.len());
+    info!("num calibration batches: {}", chunks.len());

-    info!("running onnx predictions...");
+    debug!("running onnx predictions...");
    let original_predictions = Model::run_onnx_predictions(
        &settings.run_args,
        &model_path,
@@ -970,10 +977,18 @@ pub(crate) fn calibrate(
    let pb = init_bar(range_grid.len() as u64);
    pb.set_message("calibrating...");

+    let mut num_failed = 0;
+    let mut num_passed = 0;
+
    for (((input_scale, param_scale), scale_rebase_multiplier), div_rebasing) in range_grid {
        pb.set_message(format!(
-            "input scale: {}, param scale: {}, scale rebase multiplier: {}, div rebasing: {}",
-            input_scale, param_scale, scale_rebase_multiplier, div_rebasing
+            "i-scale: {}, p-scale: {}, rebase-(x): {}, div-rebase: {}, fail: {}, pass: {}",
+            input_scale.to_string().blue(),
+            param_scale.to_string().blue(),
+            scale_rebase_multiplier.to_string().blue(),
+            div_rebasing.to_string().yellow(),
+            num_failed.to_string().red(),
+            num_passed.to_string().green()
        ));

        let key = (
@@ -1007,7 +1022,9 @@ pub(crate) fn calibrate(
        let mut circuit = match GraphCircuit::from_run_args(&local_run_args, &model_path) {
            Ok(c) => c,
            Err(e) => {
-                debug!("circuit creation from run args failed: {:?}", e);
+                error!("circuit creation from run args failed: {:?}", e);
+                pb.inc(1);
+                num_failed += 1;
                continue;
            }
        };
@@ -1039,7 +1056,9 @@ pub(crate) fn calibrate(
            Ok(_) => (),
            // typically errors will be due to the circuit overflowing the i128 limit
            Err(e) => {
-                debug!("forward pass failed: {:?}", e);
+                error!("forward pass failed: {:?}", e);
+                pb.inc(1);
+                num_failed += 1;
                continue;
            }
        }
@@ -1104,8 +1123,10 @@ pub(crate) fn calibrate(
                "found settings: \n {}",
                found_settings.as_json()?.to_colored_json_auto()?
            );
+            num_passed += 1;
        } else {
-            debug!("calibration failed {}", res.err().unwrap());
+            error!("calibration failed {}", res.err().unwrap());
+            num_failed += 1;
        }

        pb.inc(1);
@@ -1278,17 +1299,19 @@ pub(crate) fn create_evm_verifier(
    render_vk_seperately: bool,
 ) -> Result<String, Box<dyn Error>> {
    check_solc_requirement();
-    let circuit_settings = GraphSettings::load(&settings_path)?;
+
+    let settings = GraphSettings::load(&settings_path)?;
+    let commitment: Commitments = settings.run_args.commitment.into();
    let params = load_params_verifier::<KZGCommitmentScheme<Bn256>>(
        srs_path,
-        circuit_settings.run_args.logrows,
-        circuit_settings.run_args.commitment,
+        settings.run_args.logrows,
+        commitment,
    )?;

-    let num_instance = circuit_settings.total_instances();
+    let num_instance = settings.total_instances();
    let num_instance: usize = num_instance.iter().sum::<usize>();

-    let vk = load_vk::<KZGCommitmentScheme<Bn256>, GraphCircuit>(vk_path, circuit_settings)?;
+    let vk = load_vk::<KZGCommitmentScheme<Bn256>, GraphCircuit>(vk_path, settings)?;
    trace!("params computed");

    let generator = halo2_solidity_verifier::SolidityGenerator::new(
@@ -1322,17 +1345,18 @@ pub(crate) fn create_evm_vk(
    abi_path: PathBuf,
 ) -> Result<String, Box<dyn Error>> {
    check_solc_requirement();
-    let circuit_settings = GraphSettings::load(&settings_path)?;
+    let settings = GraphSettings::load(&settings_path)?;
+    let commitment: Commitments = settings.run_args.commitment.into();
    let params = load_params_verifier::<KZGCommitmentScheme<Bn256>>(
        srs_path,
-        circuit_settings.run_args.logrows,
-        circuit_settings.run_args.commitment,
+        settings.run_args.logrows,
+        commitment,
    )?;

-    let num_instance = circuit_settings.total_instances();
+    let num_instance = settings.total_instances();
    let num_instance: usize = num_instance.iter().sum::<usize>();

-    let vk = load_vk::<KZGCommitmentScheme<Bn256>, GraphCircuit>(vk_path, circuit_settings)?;
+    let vk = load_vk::<KZGCommitmentScheme<Bn256>, GraphCircuit>(vk_path, settings)?;
    trace!("params computed");

    let generator = halo2_solidity_verifier::SolidityGenerator::new(
@@ -1601,8 +1625,9 @@ pub(crate) fn setup(
    }

    let logrows = circuit.settings().run_args.logrows;
+    let commitment: Commitments = circuit.settings().run_args.commitment.into();

-    let pk = match circuit.settings().run_args.commitment {
+    let pk = match commitment {
        Commitments::KZG => {
            let params = load_params_prover::<KZGCommitmentScheme<Bn256>>(
                srs_path,
@@ -1711,7 +1736,8 @@ pub(crate) fn prove(
    let transcript: TranscriptType = proof_type.into();
    let proof_split_commits: Option<ProofSplitCommit> = data.into();

-    let commitment = circuit_settings.run_args.commitment;
+    let commitment = circuit_settings.run_args.commitment.into();
+    let logrows = circuit_settings.run_args.logrows;
    // creates and verifies the proof
    let mut snark = match commitment {
        Commitments::KZG => {
@@ -1720,7 +1746,7 @@ pub(crate) fn prove(

            let params = load_params_prover::<KZGCommitmentScheme<Bn256>>(
                srs_path,
-                circuit_settings.run_args.logrows,
+                logrows,
                Commitments::KZG,
            )?;
            match strategy {
@@ -1879,7 +1905,9 @@ pub(crate) fn mock_aggregate(
            }
            Err(_) => {
                return Err(
-                    format!("invalid sample commitment type for aggregation, must be KZG").into(),
+                    "invalid sample commitment type for aggregation, must be KZG"
+                        .to_string()
+                        .into(),
                );
            }
        }
@@ -1922,7 +1950,9 @@ pub(crate) fn setup_aggregate(
            }
            Err(_) => {
                return Err(
-                    format!("invalid sample commitment type for aggregation, must be KZG",).into(),
+                    "invalid sample commitment type for aggregation, must be KZG"
+                        .to_string()
+                        .into(),
                );
            }
        }
@@ -1983,7 +2013,9 @@ pub(crate) fn aggregate(
            }
            Err(_) => {
                return Err(
-                    format!("invalid sample commitment type for aggregation, must be KZG").into(),
+                    "invalid sample commitment type for aggregation, must be KZG"
+                        .to_string()
+                        .into(),
                );
            }
        }
@@ -2156,8 +2188,9 @@ pub(crate) fn verify(
    let circuit_settings = GraphSettings::load(&settings_path)?;

    let logrows = circuit_settings.run_args.logrows;
+    let commitment = circuit_settings.run_args.commitment.into();

-    match circuit_settings.run_args.commitment {
+    match commitment {
        Commitments::KZG => {
            let proof = Snark::load::<KZGCommitmentScheme<Bn256>>(&proof_path)?;
            let params: ParamsKZG<Bn256> = if reduced_srs {
--- a/src/graph/mod.rs
+++ b/src/graph/mod.rs
@@ -26,6 +26,7 @@ use self::input::{FileSource, GraphData};
 use self::modules::{GraphModules, ModuleConfigs, ModuleForwardResult, ModuleSizes};
 use crate::circuit::lookup::LookupOp;
 use crate::circuit::modules::ModulePlanner;
+use crate::circuit::region::ConstantsMap;
 use crate::circuit::table::{num_cols_required, Range, Table, RESERVED_BLINDING_ROWS_PAD};
 use crate::circuit::{CheckMode, InputType};
 use crate::fieldutils::felt_to_f64;
@@ -38,7 +39,7 @@ use halo2_proofs::{
    plonk::{Circuit, ConstraintSystem, Error as PlonkError},
 };
 use halo2curves::bn256::{self, Fr as Fp, G1Affine};
-use halo2curves::ff::PrimeField;
+use halo2curves::ff::{Field, PrimeField};
 #[cfg(not(target_arch = "wasm32"))]
 use lazy_static::lazy_static;
 use log::{debug, error, trace, warn};
@@ -155,7 +156,7 @@ use std::cell::RefCell;
 thread_local!(
    /// This is a global variable that holds the settings for the graph
    /// This is used to pass settings to the layouter and other parts of the circuit without needing to heavily modify the Halo2 API in a new fork
-    pub static GLOBAL_SETTINGS: RefCell<Option<GraphSettings>> = RefCell::new(None)
+    pub static GLOBAL_SETTINGS: RefCell<Option<GraphSettings>> = const { RefCell::new(None) }
 );

 /// Result from a forward pass
@@ -1051,12 +1052,10 @@ impl GraphCircuit {
    }

    fn calc_safe_lookup_range(min_max_lookup: Range, lookup_safety_margin: i128) -> Range {
-        let margin = (
+        (
            lookup_safety_margin * min_max_lookup.0,
            lookup_safety_margin * min_max_lookup.1,
-        );
-
-        margin
+        )
    }

    fn calc_num_cols(range_len: i128, max_logrows: u32) -> usize {
@@ -1240,7 +1239,7 @@ impl GraphCircuit {
        inputs: &mut [Tensor<Fp>],
        vk: Option<&VerifyingKey<G1Affine>>,
        srs: Option<&Scheme::ParamsProver>,
-        throw_range_check_error: bool,
+        witness_gen: bool,
    ) -> Result<GraphWitness, Box<dyn std::error::Error>> {
        let original_inputs = inputs.to_vec();

@@ -1289,7 +1288,7 @@ impl GraphCircuit {

        let mut model_results =
            self.model()
-                .forward(inputs, &self.settings().run_args, throw_range_check_error)?;
+                .forward(inputs, &self.settings().run_args, witness_gen)?;

        if visibility.output.requires_processing() {
            let module_outlets = visibility.output.overwrites_inputs();
@@ -1452,7 +1451,8 @@ impl GraphCircuit {
 }

 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
-struct CircuitSize {
+/// The configuration for the graph circuit
+pub struct CircuitSize {
    num_instances: usize,
    num_advice_columns: usize,
    num_fixed: usize,
@@ -1462,7 +1462,8 @@ struct CircuitSize {
 }

 impl CircuitSize {
-    pub fn from_cs(cs: &ConstraintSystem<Fp>, logrows: u32) -> Self {
+    ///
+    pub fn from_cs<F: Field>(cs: &ConstraintSystem<F>, logrows: u32) -> Self {
        CircuitSize {
            num_instances: cs.num_instance_columns(),
            num_advice_columns: cs.num_advice_columns(),
@@ -1604,6 +1605,8 @@ impl Circuit<Fp> for GraphCircuit {
        let output_vis = &self.settings().run_args.output_visibility;
        let mut graph_modules = GraphModules::new();

+        let mut constants = ConstantsMap::new();
+
        let mut config = config.clone();

        let mut inputs = self
@@ -1649,6 +1652,7 @@ impl Circuit<Fp> for GraphCircuit {
                &mut input_outlets,
                input_visibility,
                &mut instance_offset,
+                &mut constants,
            )?;
            // replace inputs with the outlets
            for (i, outlet) in outlets.iter().enumerate() {
@@ -1661,6 +1665,7 @@ impl Circuit<Fp> for GraphCircuit {
                &mut inputs,
                input_visibility,
                &mut instance_offset,
+                &mut constants,
            )?;
        }

@@ -1697,6 +1702,7 @@ impl Circuit<Fp> for GraphCircuit {
                &mut flattened_params,
                param_visibility,
                &mut instance_offset,
+                &mut constants,
            )?;

            let shapes = self.model().const_shapes();
@@ -1725,6 +1731,7 @@ impl Circuit<Fp> for GraphCircuit {
                &inputs,
                &mut vars,
                &outputs,
+                &mut constants,
            )
            .map_err(|e| {
                log::error!("{}", e);
@@ -1749,6 +1756,7 @@ impl Circuit<Fp> for GraphCircuit {
                &mut output_outlets,
                &self.settings().run_args.output_visibility,
                &mut instance_offset,
+                &mut constants,
            )?;

            // replace outputs with the outlets
@@ -1762,6 +1770,7 @@ impl Circuit<Fp> for GraphCircuit {
                &mut outputs,
                &self.settings().run_args.output_visibility,
                &mut instance_offset,
+                &mut constants,
            )?;
        }

--- a/src/graph/model.rs
+++ b/src/graph/model.rs
@@ -5,6 +5,7 @@ use super::vars::*;
 use super::GraphError;
 use super::GraphSettings;
 use crate::circuit::hybrid::HybridOp;
+use crate::circuit::region::ConstantsMap;
 use crate::circuit::region::RegionCtx;
 use crate::circuit::table::Range;
 use crate::circuit::Input;
@@ -404,7 +405,7 @@ impl ParsedNodes {
                .get(input)
                .ok_or(GraphError::MissingNode(*input))?;
            let input_dims = node.out_dims();
-            let input_dim = input_dims.get(0).ok_or(GraphError::MissingNode(*input))?;
+            let input_dim = input_dims.first().ok_or(GraphError::MissingNode(*input))?;
            inputs.push(input_dim.clone());
        }

@@ -514,21 +515,24 @@ impl Model {
            instance_shapes.len().to_string().blue(),
            "instances".blue()
        );
-        // this is the total number of variables we will need to allocate
-        // for the circuit
-        let default_value = if !self.visibility.input.is_fixed() {
-            ValType::Value(Value::<Fp>::unknown())
-        } else {
-            ValType::Constant(Fp::ONE)
-        };

        let inputs: Vec<ValTensor<Fp>> = self
            .graph
            .input_shapes()?
            .iter()
            .map(|shape| {
-                let mut t: ValTensor<Fp> =
-                    vec![default_value.clone(); shape.iter().product()].into();
+                let len = shape.iter().product();
+                let mut t: ValTensor<Fp> = (0..len)
+                    .map(|_| {
+                        if !self.visibility.input.is_fixed() {
+                            ValType::Value(Value::<Fp>::unknown())
+                        } else {
+                            ValType::Constant(Fp::random(&mut rand::thread_rng()))
+                        }
+                    })
+                    .collect::<Vec<_>>()
+                    .into();
+
                t.reshape(shape)?;
                Ok(t)
            })
@@ -577,13 +581,13 @@ impl Model {
        &self,
        model_inputs: &[Tensor<Fp>],
        run_args: &RunArgs,
-        throw_range_check_error: bool,
+        witness_gen: bool,
    ) -> Result<ForwardResult, Box<dyn Error>> {
        let valtensor_inputs: Vec<ValTensor<Fp>> = model_inputs
            .iter()
            .map(|x| x.map(|elem| ValType::Value(Value::known(elem))).into())
            .collect();
-        let res = self.dummy_layout(run_args, &valtensor_inputs, throw_range_check_error)?;
+        let res = self.dummy_layout(run_args, &valtensor_inputs, witness_gen)?;
        Ok(res.into())
    }

@@ -799,13 +803,18 @@ impl Model {
                    let input_state_idx = input_state_idx(&input_mappings);

                    let mut output_mappings = vec![];
-                    for mapping in b.output_mapping.iter() {
+                    for (i, mapping) in b.output_mapping.iter().enumerate() {
                        let mut mappings = vec![];
                        if let Some(outlet) = mapping.last_value_slot {
                            mappings.push(OutputMapping::Single {
                                outlet,
                                is_state: mapping.state,
                            });
+                        } else if mapping.state {
+                            mappings.push(OutputMapping::Single {
+                                outlet: i,
+                                is_state: mapping.state,
+                            });
                        }
                        if let Some(last) = mapping.scan {
                            mappings.push(OutputMapping::Stacked {
@@ -814,6 +823,7 @@ impl Model {
                                is_state: false,
                            });
                        }
+
                        output_mappings.push(mappings);
                    }

@@ -1071,6 +1081,8 @@ impl Model {
    /// * `layouter` - Halo2 Layouter.
    /// * `inputs` - The values to feed into the circuit.
    /// * `vars` - The variables for the circuit.
+    /// * `witnessed_outputs` - The values to compare against.
+    /// * `constants` - The constants for the circuit.
    pub fn layout(
        &self,
        mut config: ModelConfig,
@@ -1079,6 +1091,7 @@ impl Model {
        inputs: &[ValTensor<Fp>],
        vars: &mut ModelVars<Fp>,
        witnessed_outputs: &[ValTensor<Fp>],
+        constants: &mut ConstantsMap<Fp>,
    ) -> Result<Vec<ValTensor<Fp>>, Box<dyn Error>> {
        info!("model layout...");

@@ -1104,14 +1117,12 @@ impl Model {
        config.base.layout_tables(layouter)?;
        config.base.layout_range_checks(layouter)?;

-        let mut num_rows = 0;
-        let mut linear_coord = 0;
-        let mut total_const_size = 0;
+        let original_constants = constants.clone();

        let outputs = layouter.assign_region(
            || "model",
            |region| {
-                let mut thread_safe_region = RegionCtx::new(region, 0, run_args.num_inner_cols);
+                let mut thread_safe_region = RegionCtx::new_with_constants(region, 0, run_args.num_inner_cols, original_constants.clone());
                // we need to do this as this loop is called multiple times
                vars.set_instance_idx(instance_idx);

@@ -1157,29 +1168,17 @@ impl Model {
                        error!("{}", e);
                        halo2_proofs::plonk::Error::Synthesis
                    })?;
-                } else if !run_args.output_visibility.is_private() {
-                    for output in &outputs {
-                    thread_safe_region.increment_total_constants(output.num_constants());
-                    }
                }
-                num_rows = thread_safe_region.row();
-                linear_coord = thread_safe_region.linear_coord();
-                total_const_size = thread_safe_region.total_constants();
+                // Then number of columns in the circuits
+                #[cfg(not(target_arch = "wasm32"))]
+                thread_safe_region.debug_report();
+
+                *constants = thread_safe_region.assigned_constants().clone();

                Ok(outputs)
            },
-        )?;

-        // Then number of columns in the circuits
-        #[cfg(not(target_arch = "wasm32"))]
-        debug!(
-            "{} {} {} (coord={}, constants={})",
-            "model uses".blue(),
-            num_rows.to_string().blue(),
-            "rows".blue(),
-            linear_coord.to_string().yellow(),
-            total_const_size.to_string().red()
-        );
+        )?;

        let duration = start_time.elapsed();
        trace!("model layout took: {:?}", duration);
@@ -1201,6 +1200,20 @@ impl Model {
            .collect();

        for (idx, node) in self.graph.nodes.iter() {
+            debug!("laying out {}: {}", idx, node.as_str(),);
+            // Then number of columns in the circuits
+            #[cfg(not(target_arch = "wasm32"))]
+            region.debug_report();
+            debug!("input indices: {:?}", node.inputs());
+            debug!("output scales: {:?}", node.out_scales());
+            debug!(
+                "input scales: {:?}",
+                node.inputs()
+                    .iter()
+                    .map(|(idx, outlet)| self.graph.nodes[idx].out_scales()[*outlet])
+                    .collect_vec()
+            );
+
            let mut values: Vec<ValTensor<Fp>> = if !node.is_input() {
                node.inputs()
                    .iter()
@@ -1212,31 +1225,11 @@ impl Model {
                // we re-assign inputs, always from the 0 outlet
                vec![results.get(idx).ok_or(GraphError::MissingResults)?[0].clone()]
            };
-
+            debug!("output dims: {:?}", node.out_dims());
            debug!(
-                "laying out {}: {}, row:{}, coord:{}, total_constants: {}, max_lookup_inputs: {}, min_lookup_inputs: {}",
-                idx,
-                node.as_str(),
-                region.row(),
-                region.linear_coord(),
-                region.total_constants(),
-                region.max_lookup_inputs(),
-                region.min_lookup_inputs()
-            );
-            debug!("dims: {:?}", node.out_dims());
-            debug!(
-                "input_dims {:?}",
+                "input dims {:?}",
                values.iter().map(|v| v.dims()).collect_vec()
            );
-            debug!("output scales: {:?}", node.out_scales());
-            debug!("input indices: {:?}", node.inputs());
-            debug!(
-                "input scales: {:?}",
-                node.inputs()
-                    .iter()
-                    .map(|(idx, outlet)| self.graph.nodes[idx].out_scales()[*outlet])
-                    .collect_vec()
-            );

            match &node {
                NodeType::Node(n) => {
@@ -1277,8 +1270,8 @@ impl Model {
                    let num_iter = number_of_iterations(&input_mappings, input_dims.collect());

                    debug!(
-                        "{} iteration(s) in a subgraph with inputs {:?} and sources {:?}",
-                        num_iter, inputs, model.graph.inputs
+                        "{} iteration(s) in a subgraph with inputs {:?}, sources {:?}, and outputs {:?}",
+                        num_iter, inputs, model.graph.inputs, model.graph.outputs
                    );

                    let mut full_results: Vec<ValTensor<Fp>> = vec![];
@@ -1310,6 +1303,7 @@ impl Model {
                        let res = model.layout_nodes(config, region, &mut subgraph_results)?;

                        let mut outlets = BTreeMap::new();
+                        let mut stacked_outlets = BTreeMap::new();

                        for (mappings, outlet_res) in output_mappings.iter().zip(res) {
                            for mapping in mappings {
@@ -1322,25 +1316,42 @@ impl Model {
                                            let stacked_res = full_results[*outlet]
                                                .clone()
                                                .concat_axis(outlet_res.clone(), axis)?;
-
-                                            outlets.insert(outlet, stacked_res);
-                                        } else {
-                                            outlets.insert(outlet, outlet_res.clone());
+                                            stacked_outlets.insert(outlet, stacked_res);
                                        }
+                                        outlets.insert(outlet, outlet_res.clone());
                                    }
                                }
                            }
                        }

-                        full_results = outlets.into_values().collect_vec();
+                        // now extend with stacked elements
+                        let mut pre_stacked_outlets = outlets.clone();
+                        pre_stacked_outlets.extend(stacked_outlets);
+
+                        let outlets = outlets.into_values().collect_vec();
+
+                        full_results = pre_stacked_outlets.into_values().collect_vec();

                        let output_states = output_state_idx(output_mappings);
                        let input_states = input_state_idx(&input_mappings);

-                        assert_eq!(input_states.len(), output_states.len());
+                        assert_eq!(
+                            input_states.len(),
+                            output_states.len(),
+                            "input and output states must be the same length, got {:?} and {:?}",
+                            input_mappings,
+                            output_mappings
+                        );

                        for (input_idx, output_idx) in input_states.iter().zip(output_states) {
-                            values[*input_idx] = full_results[output_idx].clone();
+                            assert_eq!(
+                                values[*input_idx].dims(),
+                                outlets[output_idx].dims(),
+                                "input and output dims must be the same, got {:?} and {:?}",
+                                values[*input_idx].dims(),
+                                outlets[output_idx].dims()
+                            );
+                            values[*input_idx] = outlets[output_idx].clone();
                        }
                    }

@@ -1380,7 +1391,7 @@ impl Model {
        &self,
        run_args: &RunArgs,
        inputs: &[ValTensor<Fp>],
-        throw_range_check_error: bool,
+        witness_gen: bool,
    ) -> Result<DummyPassRes, Box<dyn Error>> {
        debug!("calculating num of constraints using dummy model layout...");

@@ -1399,29 +1410,31 @@ impl Model {
            vars: ModelVars::new_dummy(),
        };

-        let mut region = RegionCtx::new_dummy(0, run_args.num_inner_cols, throw_range_check_error);
+        let mut region = RegionCtx::new_dummy(0, run_args.num_inner_cols, witness_gen);

        let outputs = self.layout_nodes(&mut model_config, &mut region, &mut results)?;

        if self.visibility.output.is_public() || self.visibility.output.is_fixed() {
-            let default_value = if !self.visibility.output.is_fixed() {
-                ValType::Value(Value::<Fp>::unknown())
-            } else {
-                ValType::Constant(Fp::ONE)
-            };
-
            let output_scales = self.graph.get_output_scales()?;
            let res = outputs
                .iter()
                .enumerate()
                .map(|(i, output)| {
+                    let mut comparator: ValTensor<Fp> = (0..output.len())
+                        .map(|_| {
+                            if !self.visibility.output.is_fixed() {
+                                ValType::Value(Value::<Fp>::unknown())
+                            } else {
+                                ValType::Constant(Fp::random(&mut rand::thread_rng()))
+                            }
+                        })
+                        .collect::<Vec<_>>()
+                        .into();
+                    comparator.reshape(output.dims())?;
+
                    let mut tolerance = run_args.tolerance;
                    tolerance.scale = scale_to_multiplier(output_scales[i]).into();

-                    let mut comparator: ValTensor<Fp> =
-                        vec![default_value.clone(); output.dims().iter().product::<usize>()].into();
-                    comparator.reshape(output.dims())?;
-
                    dummy_config.layout(
                        &mut region,
                        &[output.clone(), comparator],
@@ -1432,7 +1445,7 @@ impl Model {
            res?;
        } else if !self.visibility.output.is_private() {
            for output in &outputs {
-                region.increment_total_constants(output.num_constants());
+                region.update_constants(output.create_constants_map());
            }
        }

@@ -1441,14 +1454,7 @@ impl Model {

        // Then number of columns in the circuits
        #[cfg(not(target_arch = "wasm32"))]
-        debug!(
-            "{} {} {} (coord={}, constants={})",
-            "model uses".blue(),
-            region.row().to_string().blue(),
-            "rows".blue(),
-            region.linear_coord().to_string().yellow(),
-            region.total_constants().to_string().red()
-        );
+        region.debug_report();

        let outputs = outputs
            .iter()
--- a/src/graph/modules.rs
+++ b/src/graph/modules.rs
@@ -2,6 +2,7 @@ use crate::circuit::modules::polycommit::{PolyCommitChip, PolyCommitConfig};
 use crate::circuit::modules::poseidon::spec::{PoseidonSpec, POSEIDON_RATE, POSEIDON_WIDTH};
 use crate::circuit::modules::poseidon::{PoseidonChip, PoseidonConfig};
 use crate::circuit::modules::Module;
+use crate::circuit::region::ConstantsMap;
 use crate::tensor::{Tensor, ValTensor};
 use halo2_proofs::circuit::Layouter;
 use halo2_proofs::plonk::{Column, ConstraintSystem, Error, Instance, VerifyingKey};
@@ -211,12 +212,13 @@ impl GraphModules {
        layouter: &mut impl Layouter<Fp>,
        x: &mut Vec<ValTensor<Fp>>,
        instance_offset: &mut usize,
+        constants: &mut ConstantsMap<Fp>,
    ) -> Result<(), Error> {
        // reserve module 0 for ... modules
        // hash the input and replace the constrained cells in the input
        let cloned_x = (*x).clone();
        x[0] = module
-            .layout(layouter, &cloned_x, instance_offset.to_owned())
+            .layout(layouter, &cloned_x, instance_offset.to_owned(), constants)
            .unwrap();
        for inc in module.instance_increment_input().iter() {
            // increment the instance offset to make way for future module layouts
@@ -234,6 +236,7 @@ impl GraphModules {
        values: &mut [ValTensor<Fp>],
        element_visibility: &Visibility,
        instance_offset: &mut usize,
+        constants: &mut ConstantsMap<Fp>,
    ) -> Result<(), Error> {
        if element_visibility.is_polycommit() && !values.is_empty() {
            // concat values and sk to get the inputs
@@ -248,7 +251,7 @@ impl GraphModules {
                layouter
                    .assign_region(|| format!("_enter_module_{}", module_offset), |_| Ok(()))
                    .unwrap();
-                Self::layout_module(&chip, layouter, x, instance_offset).unwrap();
+                Self::layout_module(&chip, layouter, x, instance_offset, constants).unwrap();
                // increment the current index
                self.polycommit_idx += 1;
            });
@@ -270,7 +273,7 @@ impl GraphModules {
                let mut inputs = values.iter_mut().map(|x| vec![x.clone()]).collect_vec();
                // layout the module
                inputs.iter_mut().for_each(|x| {
-                    Self::layout_module(&chip, layouter, x, instance_offset).unwrap();
+                    Self::layout_module(&chip, layouter, x, instance_offset, constants).unwrap();
                });
                // replace the inputs with the outputs
                values.iter_mut().enumerate().for_each(|(i, x)| {
--- a/src/graph/utilities.rs
+++ b/src/graph/utilities.rs
@@ -509,7 +509,7 @@ pub fn new_op_from_onnx(
            // if param_visibility.is_public() {
            if let Some(c) = inputs[1].opkind().get_mutable_constant() {
                inputs[1].decrement_use();
-                deleted_indices.push(inputs.len() - 1);
+                deleted_indices.push(1);
                op = SupportedOp::Linear(crate::circuit::ops::poly::PolyOp::ScatterND {
                    constant_idx: Some(c.raw_values.map(|x| x as usize)),
                })
@@ -545,7 +545,7 @@ pub fn new_op_from_onnx(
            // if param_visibility.is_public() {
            if let Some(c) = inputs[1].opkind().get_mutable_constant() {
                inputs[1].decrement_use();
-                deleted_indices.push(inputs.len() - 1);
+                deleted_indices.push(1);
                op = SupportedOp::Linear(crate::circuit::ops::poly::PolyOp::GatherND {
                    batch_dims,
                    indices: Some(c.raw_values.map(|x| x as usize)),
@@ -582,7 +582,7 @@ pub fn new_op_from_onnx(
            // if param_visibility.is_public() {
            if let Some(c) = inputs[1].opkind().get_mutable_constant() {
                inputs[1].decrement_use();
-                deleted_indices.push(inputs.len() - 1);
+                deleted_indices.push(1);
                op = SupportedOp::Linear(crate::circuit::ops::poly::PolyOp::GatherElements {
                    dim: axis,
                    constant_idx: Some(c.raw_values.map(|x| x as usize)),
@@ -734,6 +734,19 @@ pub fn new_op_from_onnx(

            SupportedOp::Linear(PolyOp::Sum { axes })
        }
+        "Reduce<MeanOfSquares>" => {
+            if inputs.len() != 1 {
+                return Err(Box::new(GraphError::InvalidDims(
+                    idx,
+                    "mean of squares".to_string(),
+                )));
+            };
+            let op = load_op::<Reduce>(node.op(), idx, node.op().name().to_string())?;
+            let axes = op.axes.into_iter().collect();
+
+            SupportedOp::Linear(PolyOp::MeanOfSquares { axes })
+        }
+
        "Max" => {
            // Extract the max value
            // first find the input that is a constant
@@ -1072,8 +1085,12 @@ pub fn new_op_from_onnx(
                }
            };

+            let in_scale = inputs[0].out_scales()[0];
+            let max_scale = std::cmp::max(scales.get_max(), in_scale);
+
            SupportedOp::Hybrid(HybridOp::Softmax {
-                scale: scale_to_multiplier(inputs[0].out_scales()[0]).into(),
+                input_scale: scale_to_multiplier(in_scale).into(),
+                output_scale: scale_to_multiplier(max_scale).into(),
                axes: softmax_op.axes.to_vec(),
            })
        }
@@ -1161,7 +1178,7 @@ pub fn new_op_from_onnx(
            // if param_visibility.is_public() {
            if let Some(c) = inputs[1].opkind().get_mutable_constant() {
                inputs[1].decrement_use();
-                deleted_indices.push(inputs.len() - 1);
+                deleted_indices.push(1);
                if c.raw_values.len() > 1 {
                    unimplemented!("only support scalar pow")
                }
--- a/src/graph/vars.rs
+++ b/src/graph/vars.rs
@@ -346,7 +346,7 @@ pub struct ModelVars<F: PrimeField + TensorType + PartialOrd> {
    pub instance: Option<ValTensor<F>>,
 }

-impl<F: PrimeField + TensorType + PartialOrd> ModelVars<F> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> ModelVars<F> {
    /// Get instance col
    pub fn get_instance_col(&self) -> Option<&Column<Instance>> {
        if let Some(instance) = &self.instance {
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -23,7 +23,6 @@
 )]
 // we allow this for our dynamic range based indexing scheme
 #![allow(clippy::single_range_in_vec_init)]
-#![feature(round_ties_even)]

 //! A library for turning computational graphs, such as neural networks, into ZK-circuits.
 //!
@@ -115,6 +114,12 @@ pub enum Commitments {
    IPA,
 }

+impl From<Option<Commitments>> for Commitments {
+    fn from(value: Option<Commitments>) -> Self {
+        value.unwrap_or(Commitments::KZG)
+    }
+}
+
 impl FromStr for Commitments {
    type Err = String;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
@@ -215,7 +220,7 @@ pub struct RunArgs {
    pub check_mode: CheckMode,
    /// commitment scheme
    #[arg(long, default_value = "kzg")]
-    pub commitment: Commitments,
+    pub commitment: Option<Commitments>,
 }

 impl Default for RunArgs {
@@ -235,7 +240,7 @@ impl Default for RunArgs {
            div_rebasing: false,
            rebase_frac_zero_constants: false,
            check_mode: CheckMode::UNSAFE,
-            commitment: Commitments::KZG,
+            commitment: None,
        }
    }
 }
--- a/src/pfsys/evm/aggregation_kzg.rs
+++ b/src/pfsys/evm/aggregation_kzg.rs
@@ -1,4 +1,8 @@
+#[cfg(not(target_arch = "wasm32"))]
+use crate::graph::CircuitSize;
 use crate::pfsys::{Snark, SnarkWitness};
+#[cfg(not(target_arch = "wasm32"))]
+use colored_json::ToColoredJson;
 use halo2_proofs::circuit::AssignedCell;
 use halo2_proofs::plonk::{self};
 use halo2_proofs::{
@@ -16,6 +20,8 @@ use halo2_wrong_ecc::{
 use halo2curves::bn256::{Bn256, Fq, Fr, G1Affine};
 use halo2curves::ff::PrimeField;
 use itertools::Itertools;
+#[cfg(not(target_arch = "wasm32"))]
+use log::debug;
 use log::trace;
 use rand::rngs::OsRng;
 use snark_verifier::loader::native::NativeLoader;
@@ -193,6 +199,23 @@ impl AggregationConfig {
        let main_gate_config = MainGate::<F>::configure(meta);
        let range_config =
            RangeChip::<F>::configure(meta, &main_gate_config, composition_bits, overflow_bits);
+
+        #[cfg(not(target_arch = "wasm32"))]
+        {
+            let circuit_size = CircuitSize::from_cs(meta, 23);
+
+            // not wasm
+
+            debug!(
+                "circuit size: \n {}",
+                circuit_size
+                    .as_json()
+                    .unwrap()
+                    .to_colored_json_auto()
+                    .unwrap()
+            );
+        }
+
        AggregationConfig {
            main_gate_config,
            range_config,
--- a/src/python.rs
+++ b/src/python.rs
@@ -197,7 +197,7 @@ impl From<PyRunArgs> for RunArgs {
            div_rebasing: py_run_args.div_rebasing,
            rebase_frac_zero_constants: py_run_args.rebase_frac_zero_constants,
            check_mode: py_run_args.check_mode,
-            commitment: py_run_args.commitment.into(),
+            commitment: Some(py_run_args.commitment.into()),
        }
    }
 }
@@ -234,6 +234,16 @@ pub enum PyCommitments {
    IPA,
 }

+impl From<Option<Commitments>> for PyCommitments {
+    fn from(commitment: Option<Commitments>) -> Self {
+        match commitment {
+            Some(Commitments::KZG) => PyCommitments::KZG,
+            Some(Commitments::IPA) => PyCommitments::IPA,
+            None => PyCommitments::KZG,
+        }
+    }
+}
+
 impl From<PyCommitments> for Commitments {
    fn from(py_commitments: PyCommitments) -> Self {
        match py_commitments {
--- a/src/tensor/mod.rs
+++ b/src/tensor/mod.rs
@@ -378,7 +378,7 @@ impl<F: PrimeField + Clone + TensorType + PartialOrd> From<Tensor<AssignedCell<A
 {
    fn from(value: Tensor<AssignedCell<Assigned<F>, F>>) -> Tensor<Value<F>> {
        let mut output = Vec::new();
-        for (_, x) in value.iter().enumerate() {
+        for x in value.iter() {
            output.push(x.value_field().evaluate());
        }
        Tensor::new(Some(&output), value.dims()).unwrap()
@@ -434,6 +434,18 @@ impl<F: PrimeField + TensorType + Clone> From<Tensor<i128>> for Tensor<Value<F>>
    }
 }

+impl<T: Clone + TensorType + std::marker::Send + std::marker::Sync>
+    maybe_rayon::iter::FromParallelIterator<T> for Tensor<T>
+{
+    fn from_par_iter<I>(par_iter: I) -> Self
+    where
+        I: maybe_rayon::iter::IntoParallelIterator<Item = T>,
+    {
+        let inner: Vec<T> = par_iter.into_par_iter().collect();
+        Tensor::new(Some(&inner), &[inner.len()]).unwrap()
+    }
+}
+
 impl<T: Clone + TensorType + std::marker::Send + std::marker::Sync>
    maybe_rayon::iter::IntoParallelIterator for Tensor<T>
 {
@@ -922,6 +934,7 @@ impl<T: Clone + TensorType> Tensor<T> {
    pub fn move_axis(&mut self, source: usize, destination: usize) -> Result<Self, TensorError> {
        assert!(source < self.dims.len());
        assert!(destination < self.dims.len());
+
        let mut new_dims = self.dims.clone();
        new_dims.remove(source);
        new_dims.insert(destination, self.dims[source]);
@@ -953,6 +966,8 @@ impl<T: Clone + TensorType> Tensor<T> {
                    old_coord[source - 1] = *c;
                } else if (i < source && source < destination)
                    || (i < destination && source > destination)
+                    || (i > source && source > destination)
+                    || (i > destination && source < destination)
                {
                    old_coord[i] = *c;
                } else if i > source && source < destination {
@@ -965,7 +980,10 @@ impl<T: Clone + TensorType> Tensor<T> {
                    ));
                }
            }
-            output.set(&coord, self.get(&old_coord));
+
+            let value = self.get(&old_coord);
+
+            output.set(&coord, value);
        }

        Ok(output)
--- a/src/tensor/ops.rs
+++ b/src/tensor/ops.rs
@@ -1875,11 +1875,7 @@ pub fn topk<T: TensorType + PartialOrd>(
    let mut indexed_a = a.clone();
    indexed_a.flatten();

-    let mut indexed_a = a
-        .iter()
-        .enumerate()
-        .map(|(i, x)| (i, x))
-        .collect::<Vec<_>>();
+    let mut indexed_a = a.iter().enumerate().collect::<Vec<_>>();

    if largest {
        indexed_a.sort_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap());
@@ -3532,12 +3528,17 @@ pub mod nonlinearities {
    }

    /// softmax layout
-    pub fn softmax_axes(a: &Tensor<i128>, scale: f64, axes: &[usize]) -> Tensor<i128> {
+    pub fn softmax_axes(
+        a: &Tensor<i128>,
+        input_scale: f64,
+        output_scale: f64,
+        axes: &[usize],
+    ) -> Tensor<i128> {
        // we want this to be as small as possible so we set the output scale to 1
        let dims = a.dims();

        if dims.len() == 1 {
-            return softmax(a, scale);
+            return softmax(a, input_scale, output_scale);
        }

        let cartesian_coord = dims[..dims.len() - 1]
@@ -3560,7 +3561,7 @@ pub mod nonlinearities {

            let softmax_input = a.get_slice(&sum_dims).unwrap();

-            let res = softmax(&softmax_input, scale);
+            let res = softmax(&softmax_input, input_scale, output_scale);

            outputs.push(res);
        }
@@ -3587,20 +3588,25 @@ pub mod nonlinearities {
    ///     Some(&[2, 2, 3, 2, 2, 0]),
    ///     &[2, 3],
    /// ).unwrap();
-    /// let result = softmax(&x, 128.0);
+    /// let result = softmax(&x, 128.0, 128.0 * 128.0);
    /// // doubles the scale of the input
-    /// let expected = Tensor::<i128>::new(Some(&[2730, 2730, 2751, 2730, 2730, 2688]), &[2, 3]).unwrap();
+    /// let expected = Tensor::<i128>::new(Some(&[2734, 2734, 2755, 2734, 2734, 2692]), &[2, 3]).unwrap();
    /// assert_eq!(result, expected);
    /// ```
-    pub fn softmax(a: &Tensor<i128>, scale: f64) -> Tensor<i128> {
+    pub fn softmax(a: &Tensor<i128>, input_scale: f64, output_scale: f64) -> Tensor<i128> {
        // the more accurate calculation is commented out and we implement as below so it matches the steps in layout

-        let exp = exp(a, scale);
+        let exp = exp(a, input_scale);

        let sum = sum(&exp).unwrap();
-        let inv_denom = recip(&sum, scale, scale);
-
-        (exp * inv_denom).unwrap()
+        let inv_denom = recip(&sum, input_scale, output_scale);
+        let mut res = (exp * inv_denom).unwrap();
+        res = res
+            .iter()
+            .map(|x| ((*x as f64) / input_scale).round() as i128)
+            .collect();
+        res.reshape(a.dims()).unwrap();
+        res
    }

    /// Applies range_check_percent
@@ -4398,6 +4404,32 @@ pub mod nonlinearities {
        let sum = sum(a).unwrap();
        const_div(&sum, (scale * a.len()) as f64)
    }
+
+    /// Mean of squares axes
+    /// # Arguments
+    /// * `a` - Tensor
+    /// * `axis` - [usize]
+    /// # Examples
+    /// ```
+    /// use ezkl::tensor::Tensor;
+    /// use ezkl::tensor::ops::nonlinearities::mean_of_squares_axes;
+    /// let x = Tensor::<i128>::new(
+    /// Some(&[2, 15, 2, 1, 1, 0]),
+    /// &[2, 3],
+    /// ).unwrap();
+    /// let result = mean_of_squares_axes(&x, &[1]);
+    /// let expected = Tensor::<i128>::new(
+    /// Some(&[78, 1]),
+    /// &[2, 1],
+    /// ).unwrap();
+    /// assert_eq!(result, expected);
+    /// ```
+    pub fn mean_of_squares_axes(a: &Tensor<i128>, axes: &[usize]) -> Tensor<i128> {
+        let square = a.map(|a_i| a_i * a_i);
+        let sum = sum_axes(&square, axes).unwrap();
+        let denominator = a.len() / sum.len();
+        const_div(&sum, denominator as f64)
+    }
 }

 /// Ops that return the transcript i.e intermediate calcs of an op
--- a/src/tensor/val.rs
+++ b/src/tensor/val.rs
@@ -1,8 +1,12 @@
+use core::{iter::FilterMap, slice::Iter};
+
+use crate::circuit::region::ConstantsMap;
+
 use super::{
    ops::{intercalate_values, pad, resize},
    *,
 };
-use halo2_proofs::{arithmetic::Field, plonk::Instance};
+use halo2_proofs::{arithmetic::Field, circuit::Cell, plonk::Instance};

 pub(crate) fn create_constant_tensor<
    F: PrimeField + TensorType + std::marker::Send + std::marker::Sync + PartialOrd,
@@ -51,6 +55,24 @@ pub enum ValType<F: PrimeField + TensorType + std::marker::Send + std::marker::S
 }

 impl<F: PrimeField + TensorType + std::marker::Send + std::marker::Sync + PartialOrd> ValType<F> {
+    /// Returns the inner cell of the [ValType].
+    pub fn cell(&self) -> Option<Cell> {
+        match self {
+            ValType::PrevAssigned(cell) => Some(cell.cell()),
+            ValType::AssignedConstant(cell, _) => Some(cell.cell()),
+            _ => None,
+        }
+    }
+
+    /// Returns the assigned cell of the [ValType].
+    pub fn assigned_cell(&self) -> Option<AssignedCell<F, F>> {
+        match self {
+            ValType::PrevAssigned(cell) => Some(cell.clone()),
+            ValType::AssignedConstant(cell, _) => Some(cell.clone()),
+            _ => None,
+        }
+    }
+
    /// Returns true if the value is previously assigned.
    pub fn is_prev_assigned(&self) -> bool {
        matches!(
@@ -293,7 +315,7 @@ impl<F: PrimeField + TensorType + PartialOrd> From<Tensor<AssignedCell<F, F>>> f
    }
 }

-impl<F: PrimeField + TensorType + PartialOrd> ValTensor<F> {
+impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> ValTensor<F> {
    /// Allocate a new [ValTensor::Instance] from the ConstraintSystem with the given tensor `dims`, optionally enabling `equality`.
    pub fn new_instance(
        cs: &mut ConstraintSystem<F>,
@@ -428,10 +450,37 @@ impl<F: PrimeField + TensorType + PartialOrd> ValTensor<F> {
    }

    /// Returns the number of constants in the [ValTensor].
-    pub fn num_constants(&self) -> usize {
+    pub fn create_constants_map_iterator(
+        &self,
+    ) -> FilterMap<Iter<'_, ValType<F>>, fn(&ValType<F>) -> Option<(F, ValType<F>)>> {
        match self {
-            ValTensor::Value { inner, .. } => inner.iter().filter(|x| x.is_constant()).count(),
-            ValTensor::Instance { .. } => 0,
+            ValTensor::Value { inner, .. } => inner.iter().filter_map(|x| {
+                if let ValType::Constant(v) = x {
+                    Some((*v, x.clone()))
+                } else {
+                    None
+                }
+            }),
+            ValTensor::Instance { .. } => {
+                unreachable!("Instance tensors do not have constants")
+            }
+        }
+    }
+
+    /// Returns the number of constants in the [ValTensor].
+    pub fn create_constants_map(&self) -> ConstantsMap<F> {
+        match self {
+            ValTensor::Value { inner, .. } => inner
+                .par_iter()
+                .filter_map(|x| {
+                    if let ValType::Constant(v) = x {
+                        Some((*v, x.clone()))
+                    } else {
+                        None
+                    }
+                })
+                .collect(),
+            ValTensor::Instance { .. } => ConstantsMap::new(),
        }
    }

--- a/src/tensor/var.rs
+++ b/src/tensor/var.rs
@@ -2,7 +2,7 @@ use std::collections::HashSet;

 use log::{debug, error, warn};

-use crate::circuit::CheckMode;
+use crate::circuit::{region::ConstantsMap, CheckMode};

 use super::*;
 /// A wrapper around Halo2's `Column<Fixed>` or `Column<Advice>`.
@@ -289,9 +289,10 @@ impl VarTensor {
        &self,
        region: &mut Region<F>,
        offset: usize,
+        coord: usize,
        constant: F,
    ) -> Result<AssignedCell<F, F>, halo2_proofs::plonk::Error> {
-        let (x, y, z) = self.cartesian_coord(offset);
+        let (x, y, z) = self.cartesian_coord(offset + coord);
        match &self {
            VarTensor::Advice { inner: advices, .. } => {
                region.assign_advice_from_constant(|| "constant", advices[x][y], z, constant)
@@ -304,33 +305,28 @@ impl VarTensor {
    }

    /// Assigns [ValTensor] to the columns of the inner tensor.
-    pub fn assign_with_omissions<F: PrimeField + TensorType + PartialOrd>(
+    pub fn assign_with_omissions<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
        &self,
        region: &mut Region<F>,
        offset: usize,
        values: &ValTensor<F>,
        omissions: &HashSet<&usize>,
+        constants: &mut ConstantsMap<F>,
    ) -> Result<ValTensor<F>, halo2_proofs::plonk::Error> {
        let mut assigned_coord = 0;
        let mut res: ValTensor<F> = match values {
            ValTensor::Instance { .. } => {
                unimplemented!("cannot assign instance to advice columns with omissions")
            }
-            ValTensor::Value { inner: v, .. } => Ok::<_, halo2_proofs::plonk::Error>(
+            ValTensor::Value { inner: v, .. } => Ok::<ValTensor<F>, halo2_proofs::plonk::Error>(
                v.enum_map(|coord, k| {
                    if omissions.contains(&coord) {
-                        return Ok(k);
+                        return Ok::<_, halo2_proofs::plonk::Error>(k);
                    }
-                    let cell = self.assign_value(region, offset, k.clone(), assigned_coord)?;
+                    let cell =
+                        self.assign_value(region, offset, k.clone(), assigned_coord, constants)?;
                    assigned_coord += 1;
-
-                    match k {
-                        ValType::Constant(f) => Ok::<ValType<F>, halo2_proofs::plonk::Error>(
-                            ValType::AssignedConstant(cell, f),
-                        ),
-                        ValType::AssignedConstant(_, f) => Ok(ValType::AssignedConstant(cell, f)),
-                        _ => Ok(ValType::PrevAssigned(cell)),
-                    }
+                    Ok::<_, halo2_proofs::plonk::Error>(cell)
                })?
                .into(),
            ),
@@ -340,11 +336,12 @@ impl VarTensor {
    }

    /// Assigns [ValTensor] to the columns of the inner tensor.
-    pub fn assign<F: PrimeField + TensorType + PartialOrd>(
+    pub fn assign<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
        &self,
        region: &mut Region<F>,
        offset: usize,
        values: &ValTensor<F>,
+        constants: &mut ConstantsMap<F>,
    ) -> Result<ValTensor<F>, halo2_proofs::plonk::Error> {
        let mut res: ValTensor<F> = match values {
            ValTensor::Instance {
@@ -382,14 +379,7 @@ impl VarTensor {
            },
            ValTensor::Value { inner: v, .. } => Ok(v
                .enum_map(|coord, k| {
-                    let cell = self.assign_value(region, offset, k.clone(), coord)?;
-                    match k {
-                        ValType::Constant(f) => Ok::<ValType<F>, halo2_proofs::plonk::Error>(
-                            ValType::AssignedConstant(cell, f),
-                        ),
-                        ValType::AssignedConstant(_, f) => Ok(ValType::AssignedConstant(cell, f)),
-                        _ => Ok(ValType::PrevAssigned(cell)),
-                    }
+                    self.assign_value(region, offset, k.clone(), coord, constants)
                })?
                .into()),
        }?;
@@ -399,13 +389,16 @@ impl VarTensor {

    /// Assigns specific values (`ValTensor`) to the columns of the inner tensor but allows for column wrapping for accumulated operations.
    /// Duplication occurs by copying the last cell of the column to the first cell next column and creating a copy constraint between the two.
-    pub fn dummy_assign_with_duplication<F: PrimeField + TensorType + PartialOrd>(
+    pub fn dummy_assign_with_duplication<
+        F: PrimeField + TensorType + PartialOrd + std::hash::Hash,
+    >(
        &self,
        row: usize,
        offset: usize,
        values: &ValTensor<F>,
        single_inner_col: bool,
-    ) -> Result<(ValTensor<F>, usize, usize), halo2_proofs::plonk::Error> {
+        constants: &mut ConstantsMap<F>,
+    ) -> Result<(ValTensor<F>, usize), halo2_proofs::plonk::Error> {
        match values {
            ValTensor::Instance { .. } => unimplemented!("duplication is not supported on instance columns. increase K if you require more rows."),
            ValTensor::Value { inner: v, dims , ..} => {
@@ -430,21 +423,24 @@ impl VarTensor {

                // duplicates every nth element to adjust for column overflow
                let mut res: ValTensor<F> = v.duplicate_every_n(duplication_freq, num_repeats, duplication_offset).unwrap().into();
+
+                let constants_map = res.create_constants_map();
+                constants.extend(constants_map);
+
                let total_used_len = res.len();
-                let total_constants = res.num_constants();
                res.remove_every_n(duplication_freq, num_repeats, duplication_offset).unwrap();

                res.reshape(dims).unwrap();
                res.set_scale(values.scale());

-                Ok((res, total_used_len, total_constants))
+                Ok((res, total_used_len))
            }
        }
    }

    /// Assigns specific values (`ValTensor`) to the columns of the inner tensor but allows for column wrapping for accumulated operations.
    /// Duplication occurs by copying the last cell of the column to the first cell next column and creating a copy constraint between the two.
-    pub fn assign_with_duplication<F: PrimeField + TensorType + PartialOrd>(
+    pub fn assign_with_duplication<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
        &self,
        region: &mut Region<F>,
        row: usize,
@@ -452,7 +448,8 @@ impl VarTensor {
        values: &ValTensor<F>,
        check_mode: &CheckMode,
        single_inner_col: bool,
-    ) -> Result<(ValTensor<F>, usize, usize), halo2_proofs::plonk::Error> {
+        constants: &mut ConstantsMap<F>,
+    ) -> Result<(ValTensor<F>, usize), halo2_proofs::plonk::Error> {
        let mut prev_cell = None;

        match values {
@@ -494,7 +491,7 @@ impl VarTensor {
                        assert_eq!(Into::<i32>::into(k.clone()), Into::<i32>::into(v[coord - 1].clone()));
                    };

-                    let cell = self.assign_value(region, offset, k.clone(), coord * step)?;
+                    let cell = self.assign_value(region, offset, k.clone(), coord * step, constants)?;

                    if single_inner_col {
                    if z == 0 {
@@ -502,28 +499,23 @@ impl VarTensor {
                        prev_cell = Some(cell.clone());
                    } else if coord > 0 && z == 0 && single_inner_col {
                        if let Some(prev_cell) = prev_cell.as_ref() {
-                            region.constrain_equal(prev_cell.cell(),cell.cell())?;
+                            let cell = cell.cell().ok_or({
+                                error!("Error getting cell: {:?}", (x,y));
+                                halo2_proofs::plonk::Error::Synthesis})?;
+                            let prev_cell = prev_cell.cell().ok_or({
+                                error!("Error getting cell: {:?}", (x,y));
+                                halo2_proofs::plonk::Error::Synthesis})?;
+                            region.constrain_equal(prev_cell,cell)?;
                        } else {
                            error!("Error copy-constraining previous value: {:?}", (x,y));
                            return Err(halo2_proofs::plonk::Error::Synthesis);
                        }
                    }}

-                    match k {
-                        ValType::Constant(f) => {
-                            Ok(ValType::AssignedConstant(cell, f))
-                        },
-                        ValType::AssignedConstant(_, f) => {
-                            Ok(ValType::AssignedConstant(cell, f))
-                        },
-                        _ => {
-                            Ok(ValType::PrevAssigned(cell))
-                        }
-                    }
+                    Ok(cell)

                })?.into()};
                let total_used_len = res.len();
-                let total_constants = res.num_constants();
                res.remove_every_n(duplication_freq, num_repeats, duplication_offset).unwrap();

                res.reshape(dims).unwrap();
@@ -542,42 +534,61 @@ impl VarTensor {
                    )};
                }

-                Ok((res, total_used_len, total_constants))
+                Ok((res, total_used_len))
            }
        }
    }

-    fn assign_value<F: PrimeField + TensorType + PartialOrd>(
+    fn assign_value<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
        &self,
        region: &mut Region<F>,
        offset: usize,
        k: ValType<F>,
        coord: usize,
-    ) -> Result<AssignedCell<F, F>, halo2_proofs::plonk::Error> {
+        constants: &mut ConstantsMap<F>,
+    ) -> Result<ValType<F>, halo2_proofs::plonk::Error> {
        let (x, y, z) = self.cartesian_coord(offset + coord);
-        match k {
+        let res = match k {
            ValType::Value(v) => match &self {
                VarTensor::Advice { inner: advices, .. } => {
-                    region.assign_advice(|| "k", advices[x][y], z, || v)
+                    ValType::PrevAssigned(region.assign_advice(|| "k", advices[x][y], z, || v)?)
                }
                _ => unimplemented!(),
            },
-            ValType::PrevAssigned(v) | ValType::AssignedConstant(v, ..) => match &self {
+            ValType::PrevAssigned(v) => match &self {
                VarTensor::Advice { inner: advices, .. } => {
-                    v.copy_advice(|| "k", region, advices[x][y], z)
+                    ValType::PrevAssigned(v.copy_advice(|| "k", region, advices[x][y], z)?)
                }
-                _ => {
-                    error!("PrevAssigned is only supported for advice columns");
-                    Err(halo2_proofs::plonk::Error::Synthesis)
+                _ => unimplemented!(),
+            },
+            ValType::AssignedConstant(v, val) => match &self {
+                VarTensor::Advice { inner: advices, .. } => {
+                    ValType::AssignedConstant(v.copy_advice(|| "k", region, advices[x][y], z)?, val)
                }
+                _ => unimplemented!(),
            },
            ValType::AssignedValue(v) => match &self {
-                VarTensor::Advice { inner: advices, .. } => region
-                    .assign_advice(|| "k", advices[x][y], z, || v)
-                    .map(|a| a.evaluate()),
+                VarTensor::Advice { inner: advices, .. } => ValType::PrevAssigned(
+                    region
+                        .assign_advice(|| "k", advices[x][y], z, || v)?
+                        .evaluate(),
+                ),
                _ => unimplemented!(),
            },
-            ValType::Constant(v) => self.assign_constant(region, offset + coord, v),
-        }
+            ValType::Constant(v) => {
+                if let std::collections::hash_map::Entry::Vacant(e) = constants.entry(v) {
+                    let value = ValType::AssignedConstant(
+                        self.assign_constant(region, offset, coord, v)?,
+                        v,
+                    );
+                    e.insert(value.clone());
+                    value
+                } else {
+                    let cell = constants.get(&v).unwrap();
+                    self.assign_value(region, offset, cell.clone(), coord, constants)?
+                }
+            }
+        };
+        Ok(res)
    }
 }
--- a/src/wasm.rs
+++ b/src/wasm.rs
@@ -8,12 +8,14 @@ use crate::graph::quantize_float;
 use crate::graph::scale_to_multiplier;
 use crate::graph::{GraphCircuit, GraphSettings};
 use crate::pfsys::create_proof_circuit;
+use crate::pfsys::evm::aggregation_kzg::AggregationCircuit;
 use crate::pfsys::evm::aggregation_kzg::PoseidonTranscript;
 use crate::pfsys::verify_proof_circuit;
 use crate::pfsys::TranscriptType;
 use crate::tensor::TensorType;
 use crate::CheckMode;
 use crate::Commitments;
+use console_error_panic_hook;
 use halo2_proofs::plonk::*;
 use halo2_proofs::poly::commitment::{CommitmentScheme, ParamsProver};
 use halo2_proofs::poly::ipa::multiopen::{ProverIPA, VerifierIPA};
@@ -33,11 +35,10 @@ use halo2curves::bn256::{Bn256, Fr, G1Affine};
 use halo2curves::ff::{FromUniformBytes, PrimeField};
 use snark_verifier::loader::native::NativeLoader;
 use snark_verifier::system::halo2::transcript::evm::EvmTranscript;
+use std::str::FromStr;
 use wasm_bindgen::prelude::*;
 use wasm_bindgen_console_logger::DEFAULT_LOGGER;

-use console_error_panic_hook;
-
 #[cfg(feature = "web")]
 pub use wasm_bindgen_rayon::init_thread_pool;

@@ -336,8 +337,94 @@ pub fn verify(

    let orig_n = 1 << circuit_settings.run_args.logrows;

+    let commitment = circuit_settings.run_args.commitment.into();
+
    let mut reader = std::io::BufReader::new(&srs[..]);
-    let result = match circuit_settings.run_args.commitment {
+    let result = match commitment {
+        Commitments::KZG => {
+            let params: ParamsKZG<Bn256> =
+                halo2_proofs::poly::commitment::Params::<'_, G1Affine>::read(&mut reader)
+                    .map_err(|e| JsError::new(&format!("Failed to deserialize params: {}", e)))?;
+            let strategy = KZGSingleStrategy::new(params.verifier_params());
+            match proof.transcript_type {
+                TranscriptType::EVM => verify_proof_circuit::<
+                    VerifierSHPLONK<'_, Bn256>,
+                    KZGCommitmentScheme<Bn256>,
+                    KZGSingleStrategy<_>,
+                    _,
+                    EvmTranscript<G1Affine, _, _, _>,
+                >(&proof, &params, &vk, strategy, orig_n),
+
+                TranscriptType::Poseidon => {
+                    verify_proof_circuit::<
+                        VerifierSHPLONK<'_, Bn256>,
+                        KZGCommitmentScheme<Bn256>,
+                        KZGSingleStrategy<_>,
+                        _,
+                        PoseidonTranscript<NativeLoader, _>,
+                    >(&proof, &params, &vk, strategy, orig_n)
+                }
+            }
+        }
+        Commitments::IPA => {
+            let params: ParamsIPA<_> =
+                halo2_proofs::poly::commitment::Params::<'_, G1Affine>::read(&mut reader)
+                    .map_err(|e| JsError::new(&format!("Failed to deserialize params: {}", e)))?;
+            let strategy = IPASingleStrategy::new(params.verifier_params());
+            match proof.transcript_type {
+                TranscriptType::EVM => verify_proof_circuit::<
+                    VerifierIPA<_>,
+                    IPACommitmentScheme<G1Affine>,
+                    IPASingleStrategy<_>,
+                    _,
+                    EvmTranscript<G1Affine, _, _, _>,
+                >(&proof, &params, &vk, strategy, orig_n),
+                TranscriptType::Poseidon => {
+                    verify_proof_circuit::<
+                        VerifierIPA<_>,
+                        IPACommitmentScheme<G1Affine>,
+                        IPASingleStrategy<_>,
+                        _,
+                        PoseidonTranscript<NativeLoader, _>,
+                    >(&proof, &params, &vk, strategy, orig_n)
+                }
+            }
+        }
+    };
+
+    match result {
+        Ok(_) => Ok(true),
+        Err(e) => Err(JsError::new(&format!("{}", e))),
+    }
+}
+
+#[wasm_bindgen]
+#[allow(non_snake_case)]
+/// Verify aggregate proof in browser using wasm
+pub fn verifyAggr(
+    proof_js: wasm_bindgen::Clamped<Vec<u8>>,
+    vk: wasm_bindgen::Clamped<Vec<u8>>,
+    logrows: u64,
+    srs: wasm_bindgen::Clamped<Vec<u8>>,
+    commitment: &str,
+) -> Result<bool, JsError> {
+    let proof: crate::pfsys::Snark<Fr, G1Affine> = serde_json::from_slice(&proof_js[..])
+        .map_err(|e| JsError::new(&format!("Failed to deserialize proof: {}", e)))?;
+
+    let mut reader = std::io::BufReader::new(&vk[..]);
+    let vk = VerifyingKey::<G1Affine>::read::<_, AggregationCircuit>(
+        &mut reader,
+        halo2_proofs::SerdeFormat::RawBytes,
+        (),
+    )
+    .map_err(|e| JsError::new(&format!("Failed to deserialize vk: {}", e)))?;
+
+    let commit = Commitments::from_str(commitment).map_err(|e| JsError::new(&format!("{}", e)))?;
+
+    let orig_n = 1 << logrows;
+
+    let mut reader = std::io::BufReader::new(&srs[..]);
+    let result = match commit {
        Commitments::KZG => {
            let params: ParamsKZG<Bn256> =
                halo2_proofs::poly::commitment::Params::<'_, G1Affine>::read(&mut reader)
@@ -436,8 +523,9 @@ pub fn prove(

    // read in kzg params
    let mut reader = std::io::BufReader::new(&srs[..]);
+    let commitment = circuit.settings().run_args.commitment.into();
    // creates and verifies the proof
-    let proof = match circuit.settings().run_args.commitment {
+    let proof = match commitment {
        Commitments::KZG => {
            let params: ParamsKZG<Bn256> =
                halo2_proofs::poly::commitment::Params::<'_, G1Affine>::read(&mut reader)
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -122,7 +122,7 @@ mod native_tests {
        let settings: GraphSettings = serde_json::from_str(&settings).unwrap();
        let logrows = settings.run_args.logrows;

-        download_srs(logrows, settings.run_args.commitment);
+        download_srs(logrows, settings.run_args.commitment.into());
    }

    fn mv_test_(test_dir: &str, test: &str) {
@@ -200,7 +200,7 @@ mod native_tests {
        "1l_tiny_div",
    ];

-    const TESTS: [&str; 91] = [
+    const TESTS: [&str; 92] = [
        "1l_mlp", //0
        "1l_slice",
        "1l_concat",
@@ -296,6 +296,7 @@ mod native_tests {
        "reducel1",
        "reducel2", // 89
        "1l_lppool",
+        "lstm_large", // 91
    ];

    const WASM_TESTS: [&str; 46] = [
@@ -534,7 +535,7 @@ mod native_tests {
            }
        });

-            seq!(N in 0..=90 {
+            seq!(N in 0..=91 {

            #(#[test_case(TESTS[N])])*
            #[ignore]
@@ -622,7 +623,7 @@ mod native_tests {
            #(#[test_case(TESTS[N])])*
            fn mock_large_batch_public_outputs_(test: &str) {
                // currently variable output rank is not supported in ONNX
-                if test != "gather_nd" {
+                if test != "gather_nd" && test != "lstm_large" {
                    crate::native_tests::init_binary();
                    let test_dir = TempDir::new(test).unwrap();
                    let path = test_dir.path().to_str().unwrap(); crate::native_tests::mv_test_(path, test);
@@ -1970,7 +1971,7 @@ mod native_tests {
            .expect("failed to parse settings file");

        // get_srs for the graph_settings_num_instances
-        download_srs(1, graph_settings.run_args.commitment);
+        download_srs(1, graph_settings.run_args.commitment.into());

        let status = Command::new(format!("{}/release/ezkl", *CARGO_TARGET_DIR))
            .args([
--- a/tests/py_integration_tests.rs
+++ b/tests/py_integration_tests.rs
@@ -56,33 +56,38 @@ mod py_tests {
            // source .env/bin/activate
            // pip install -r requirements.txt
            // maturin develop --release --features python-bindings
-
+            // first install tf2onnx as it has protobuf conflict with onnx
+            let status = Command::new("pip")
+                .args(["install", "tf2onnx==1.16.1"])
+                .status()
+                .expect("failed to execute process");
+            assert!(status.success());
            // now install torch, pandas, numpy, seaborn, jupyter
            let status = Command::new("pip")
                .args([
                    "install",
-                    "torch-geometric==2.5.0",
-                    "torch==2.0.1",
-                    "torchvision==0.15.2",
-                    "pandas==2.0.3",
-                    "numpy==1.23",
-                    "seaborn==0.12.2",
-                    "jupyter==1.0.0",
-                    "onnx==1.14.0",
-                    "kaggle==1.5.15",
-                    "py-solc-x==1.1.1",
-                    "web3==6.5.0",
-                    "librosa==0.10.0.post2",
-                    "keras==2.12.0",
-                    "tensorflow==2.12.0",
-                    "tensorflow-datasets==4.9.3",
-                    "tf2onnx==1.14.0",
-                    "pytorch-lightning==2.0.6",
+                    "torch-geometric==2.5.2",
+                    "torch==2.2.2",
+                    "torchvision==0.17.2",
+                    "pandas==2.2.1",
+                    "numpy==1.26.4",
+                    "seaborn==0.13.2",
+                    "notebook==7.1.2",
+                    "nbconvert==7.16.3",
+                    "onnx==1.16.0",
+                    "kaggle==1.6.8",
+                    "py-solc-x==2.0.2",
+                    "web3==6.16.0",
+                    "librosa==0.10.1",
+                    "keras==3.1.1",
+                    "tensorflow==2.16.1",
+                    "tensorflow-datasets==4.9.4",
+                    "pytorch-lightning==2.2.1",
                    "sk2torch==1.2.0",
-                    "scikit-learn==1.3.1",
-                    "xgboost==1.7.6",
-                    "hummingbird-ml==0.4.9",
-                    "lightgbm==4.0.0",
+                    "scikit-learn==1.4.1.post1",
+                    "xgboost==2.0.3",
+                    "hummingbird-ml==0.4.11",
+                    "lightgbm==4.3.0",
                ])
                .status()
                .expect("failed to execute process");
--- a/tests/wasm.rs
+++ b/tests/wasm.rs
@@ -11,7 +11,7 @@ mod wasm32 {
        bufferToVecOfFelt, compiledCircuitValidation, encodeVerifierCalldata, feltToBigEndian,
        feltToFloat, feltToInt, feltToLittleEndian, genPk, genVk, genWitness, inputValidation,
        pkValidation, poseidonHash, proofValidation, prove, settingsValidation, srsValidation,
-        u8_array_to_u128_le, verify, vkValidation, witnessValidation,
+        u8_array_to_u128_le, verify, verifyAggr, vkValidation, witnessValidation,
    };
    use halo2_solidity_verifier::encode_calldata;
    use halo2curves::bn256::{Fr, G1Affine};
@@ -27,10 +27,29 @@ mod wasm32 {
    pub const NETWORK: &[u8] = include_bytes!("../tests/wasm/network.onnx");
    pub const INPUT: &[u8] = include_bytes!("../tests/wasm/input.json");
    pub const PROOF: &[u8] = include_bytes!("../tests/wasm/proof.json");
+    pub const PROOF_AGGR: &[u8] = include_bytes!("../tests/wasm/proof_aggr.json");
    pub const SETTINGS: &[u8] = include_bytes!("../tests/wasm/settings.json");
    pub const PK: &[u8] = include_bytes!("../tests/wasm/pk.key");
    pub const VK: &[u8] = include_bytes!("../tests/wasm/vk.key");
+    pub const VK_AGGR: &[u8] = include_bytes!("../tests/wasm/vk_aggr.key");
    pub const SRS: &[u8] = include_bytes!("../tests/wasm/kzg");
+    pub const SRS1: &[u8] = include_bytes!("../tests/wasm/kzg1.srs");
+
+    #[wasm_bindgen_test]
+    async fn can_verify_aggr() {
+        let value = verifyAggr(
+            wasm_bindgen::Clamped(PROOF_AGGR.to_vec()),
+            wasm_bindgen::Clamped(VK_AGGR.to_vec()),
+            21,
+            wasm_bindgen::Clamped(SRS1.to_vec()),
+            "kzg",
+        )
+        .map_err(|_| "failed")
+        .unwrap();
+
+        // should not fail
+        assert!(value);
+    }

    #[wasm_bindgen_test]
    async fn verify_encode_verifier_calldata() {
--- a/tests/wasm/kzg1.srs
+++ b/tests/wasm/kzg1.srs
--- a/tests/wasm/model.compiled
+++ b/tests/wasm/model.compiled
--- a/tests/wasm/proof_aggr.json
+++ b/tests/wasm/proof_aggr.json
--- a/tests/wasm/settings.json
+++ b/tests/wasm/settings.json
@@ -24,8 +24,7 @@
        "param_visibility": "Private",
        "div_rebasing": false,
        "rebase_frac_zero_constants": false,
-        "check_mode": "UNSAFE",
-        "commitment": "KZG"
+        "check_mode": "UNSAFE"
    },
    "num_rows": 16,
    "total_dynamic_col_size": 0,
--- a/tests/wasm/vk_aggr.key
+++ b/tests/wasm/vk_aggr.key
Author	SHA1	Message	Date
dante	316a9a3b40	chore: update tract (#766 )	2024-04-04 18:07:08 +01:00
dante	5389012b68	fix: patch large batch ex (#763 )	2024-04-03 02:33:57 +01:00
dante	48223cca11	fix: make commitment optional for backwards compat (#762 )	2024-04-03 02:26:50 +01:00
dante	32c3a5e159	fix: hold stacked outputs in a separate map	2024-04-02 21:37:20 +01:00
dante	ff563e93a7	fix: bump python version (#761 )	2024-04-02 17:08:26 +01:00
dante	5639d36097	chore: verify aggr wasm unit test (#760 )	2024-04-01 20:54:20 +01:00
dante	4ec8d13082	chore: verify aggr in wasm (#758 )	2024-03-29 23:28:20 +00:00
dante	12735aefd4	chore: reduce softmax recip DR (#756 )	2024-03-27 01:14:29 +00:00
dante	7fe179b8d4	feat: dictionary of reusable constants (#754 )	2024-03-26 13:12:09 +00:00
Ethan Cemer	3be988a6a0	fix: use pnpm in build script for in-browser-evm-verifier (#752 )	2024-03-25 23:23:02 +00:00