ci: update version string in docs

docs: automated sphinx documentation for python bindings (#714 )
--------- Co-authored-by: dante <45801863+alexander-camuto@users.noreply.github.com> Co-authored-by: Ethan Cemer <tylercemer@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
2026-01-13 08:17:57 -05:00 · 2024-04-05 17:33:23 +00:00 · 2024-04-05 18:33:06 +01: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
47 changed files with 5209 additions and 396 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
@@ -128,6 +128,7 @@ jobs:
          mv Cargo.lock Cargo.lock.orig
          sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" Cargo.lock.orig >Cargo.lock

+
      - name: Install required libraries
        shell: bash
        run: |
@@ -176,7 +177,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'
@@ -228,7 +229,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
@@ -263,7 +264,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"

@@ -287,7 +288,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
@@ -359,3 +360,17 @@ jobs:
        with:
          repository-url: https://test.pypi.org/legacy/
          packages-dir: ./
+
+  doc-publish:
+    name: Trigger ReadTheDocs Build
+    runs-on: ubuntu-latest
+    needs: pypi-publish
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Trigger RTDs build
+        uses: dfm/rtds-action@v1
+        with:
+          webhook_url: ${{ secrets.RTDS_WEBHOOK_URL }}
+          webhook_token: ${{ secrets.RTDS_WEBHOOK_TOKEN }}
+          commit_ref: ${{ github.ref_name }}
--- 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/.github/workflows/tagging.yml
+++ b/.github/workflows/tagging.yml
@@ -14,6 +14,40 @@ jobs:
      - uses: actions/checkout@v4
      - name: Bump version and push tag
        id: tag_version
-        uses: mathieudutour/github-tag-action@v6.1
+        uses: mathieudutour/github-tag-action@v6.2
        with:
          github_token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Set Cargo.toml version to match github tag for docs
+        shell: bash
+        env:
+          RELEASE_TAG: ${{ steps.tag_version.outputs.new_tag }}
+        run: |
+          mv docs/python/src/conf.py docs/python/src/conf.py.orig
+          sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" docs/python/src/conf.py.orig >docs/python/src/conf.py
+          rm docs/python/src/conf.py.orig
+          mv docs/python/requirements-docs.txt docs/python/requirements-docs.txt.orig
+          sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" docs/python/requirements-docs.txt.orig >docs/python/requirements-docs.txt
+          rm docs/python/requirements-docs.txt.orig
+
+      - name: Commit files and create tag
+        env:
+          RELEASE_TAG: ${{ steps.tag_version.outputs.new_tag }}
+        run: |
+          git config --local user.email "github-actions[bot]@users.noreply.github.com"
+          git config --local user.name "github-actions[bot]"
+          git fetch --tags
+          git checkout -b release-$RELEASE_TAG
+          git add .
+          git commit -m "ci: update version string in docs"
+          git tag -d $RELEASE_TAG
+          git tag $RELEASE_TAG
+
+      - name: Push changes
+        uses: ad-m/github-push-action@master
+        env:
+          RELEASE_TAG: ${{ steps.tag_version.outputs.new_tag }}
+        with:
+          branch: release-${{ steps.tag_version.outputs.new_tag }}
+          force: true
+          tags: true
--- a/.gitignore
+++ b/.gitignore
@@ -48,4 +48,6 @@ node_modules
 /dist
 timingData.json
 !tests/wasm/pk.key
-!tests/wasm/vk.key
+!tests/wasm/vk.key
+docs/python/build
+!tests/wasm/vk_aggr.key
--- a/.python-version
+++ b/.python-version
@@ -0,0 +1 @@
+3.12.1
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -0,0 +1,26 @@
+# .readthedocs.yaml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+version: 2
+
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.12"
+
+# Build documentation in the "docs/" directory with Sphinx
+sphinx:
+  configuration: ./docs/python/src/conf.py
+
+# Optionally build your docs in additional formats such as PDF and ePub
+# formats:
+#    - pdf
+#    - epub
+
+# Optional but recommended, declare the Python requirements required
+# to build your documentation
+# See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
+python:
+  install:
+  - requirements: ./docs/python/requirements-docs.txt
--- 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/docs/python/build.sh
+++ b/docs/python/build.sh
@@ -0,0 +1,2 @@
+#!/bin/sh
+sphinx-build ./src build
--- a/docs/python/requirements-docs.txt
+++ b/docs/python/requirements-docs.txt
@@ -0,0 +1,4 @@
+ezkl==10.2.9
+sphinx
+sphinx-rtd-theme
+sphinxcontrib-napoleon
--- a/docs/python/src/conf.py
+++ b/docs/python/src/conf.py
@@ -0,0 +1,29 @@
+import ezkl
+
+project = 'ezkl'
+release = '10.2.9'
+version = release
+
+
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.autosummary',
+    'sphinx.ext.intersphinx',
+    'sphinx.ext.todo',
+    'sphinx.ext.inheritance_diagram',
+    'sphinx.ext.autosectionlabel',
+    'sphinx.ext.napoleon',
+    'sphinx_rtd_theme',
+]
+
+autosummary_generate = True
+autosummary_imported_members = True
+
+templates_path = ['_templates']
+exclude_patterns = []
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = 'sphinx_rtd_theme'
+html_static_path = ['_static']
--- a/docs/python/src/index.rst
+++ b/docs/python/src/index.rst
@@ -0,0 +1,11 @@
+.. extension documentation master file, created by
+   sphinx-quickstart on Mon Jun 19 15:02:05 2023.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+ezkl python bindings
+================================================
+
+.. automodule:: ezkl
+   :members:
+   :undoc-members:
--- 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/pyproject.toml
+++ b/pyproject.toml
@@ -1,5 +1,5 @@
 [build-system]
-requires = ["maturin>=0.14,<0.15"]
+requires = ["maturin>=1.0,<2.0"]
 build-backend = "maturin"

 [tool.pytest.ini_options]
--- 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.1
+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/ops/layouts.rs
+++ b/src/circuit/ops/layouts.rs
@@ -568,10 +568,10 @@ fn _sort_ascending<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
    let is_assigned = !input.any_unknowns()?;

    let sorted = if is_assigned {
-        input
-            .get_int_evals()?
-            .iter()
-            .sorted_by(|a, b| a.cmp(b))
+        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 {
@@ -753,20 +753,28 @@ pub(crate) fn dynamic_lookup<F: PrimeField + TensorType + PartialOrd + std::hash
    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| {
@@ -1967,6 +1975,21 @@ pub(crate) fn pairwise<F: PrimeField + TensorType + PartialOrd + std::hash::Hash
    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 + std::hash::Hash>(
    config: &BaseConfig<F>,
@@ -3251,11 +3274,15 @@ pub(crate) fn softmax<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>
    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,
-        values,
+        &[sub],
        &LookupOp::Exp { scale: input_scale },
    )?;

--- 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,
@@ -105,10 +108,28 @@ impl<

    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(),
@@ -146,6 +167,10 @@ impl<
    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(
@@ -292,6 +317,9 @@ impl<
            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()?)?,
@@ -404,6 +432,7 @@ impl<

    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
@@ -163,7 +163,7 @@ impl<'a, F: PrimeField + TensorType + PartialOrd + std::hash::Hash> RegionCtx<'a

    ///
    pub fn update_constants(&mut self, constants: ConstantsMap<F>) {
-        self.assigned_constants.extend(constants.into_iter());
+        self.assigned_constants.extend(constants);
    }

    ///
@@ -389,7 +389,7 @@ impl<'a, F: PrimeField + TensorType + PartialOrd + std::hash::Hash> RegionCtx<'a
                shuffle_index.update(&local_reg.shuffle_index);
                // update the constants
                let mut constants = constants.lock().unwrap();
-                constants.extend(local_reg.assigned_constants.into_iter());
+                constants.extend(local_reg.assigned_constants);

                res
            })
@@ -574,8 +574,10 @@ impl<'a, F: PrimeField + TensorType + PartialOrd + std::hash::Hash> RegionCtx<'a
                &mut self.assigned_constants,
            )
        } else {
-            let values_map = values.create_constants_map();
-            self.assigned_constants.extend(values_map);
+            if !values.is_instance() {
+                let values_map = values.create_constants_map_iterator();
+                self.assigned_constants.extend(values_map);
+            }
            Ok(values.clone())
        }
    }
@@ -599,8 +601,10 @@ impl<'a, F: PrimeField + TensorType + PartialOrd + std::hash::Hash> RegionCtx<'a
                &mut self.assigned_constants,
            )
        } else {
-            let values_map = values.create_constants_map();
-            self.assigned_constants.extend(values_map);
+            if !values.is_instance() {
+                let values_map = values.create_constants_map_iterator();
+                self.assigned_constants.extend(values_map);
+            }
            Ok(values.clone())
        }
    }
@@ -630,9 +634,8 @@ impl<'a, F: PrimeField + TensorType + PartialOrd + std::hash::Hash> RegionCtx<'a
                &mut self.assigned_constants,
            )
        } else {
-            let mut values_map = values.create_constants_map();
-
            let inner_tensor = values.get_inner_tensor().unwrap();
+            let mut values_map = values.create_constants_map();

            for o in ommissions {
                if let ValType::Constant(value) = inner_tensor.get_flat_index(**o) {
--- 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
@@ -339,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,
@@ -360,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 {
@@ -384,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"))]
@@ -586,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());
        }
@@ -666,21 +666,17 @@ 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,
@@ -694,7 +690,7 @@ 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,
@@ -911,7 +907,7 @@ 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()?)?;
-    debug!("num of calibration batches: {}", chunks.len());
+    info!("num calibration batches: {}", chunks.len());

    debug!("running onnx predictions...");
    let original_predictions = Model::run_onnx_predictions(
@@ -1303,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(
@@ -1347,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(
@@ -1626,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,
@@ -1736,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 => {
@@ -1745,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 {
@@ -2187,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
@@ -39,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};
@@ -1451,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,
@@ -1461,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(),
--- a/src/graph/model.rs
+++ b/src/graph/model.rs
@@ -803,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 {
@@ -818,6 +823,7 @@ impl Model {
                                is_state: false,
                            });
                        }
+
                        output_mappings.push(mappings);
                    }

@@ -1194,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()
@@ -1205,25 +1225,11 @@ impl Model {
                // we re-assign inputs, always from the 0 outlet
                vec![results.get(idx).ok_or(GraphError::MissingResults)?[0].clone()]
            };
-
-            debug!("laying out {}: {}", idx, node.as_str(),);
-            // Then number of columns in the circuits
-            #[cfg(not(target_arch = "wasm32"))]
-            region.debug_report();
-            debug!("dims: {:?}", node.out_dims());
+            debug!("output 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) => {
@@ -1264,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![];
@@ -1297,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 {
@@ -1309,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();
                        }
                    }

--- 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
@@ -1165,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/lib.rs
+++ b/src/lib.rs
@@ -114,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> {
@@ -214,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 {
@@ -234,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
@@ -33,6 +33,7 @@ use tokio::runtime::Runtime;

 type PyFelt = String;

+/// pyclass representing an enum
 #[pyclass]
 #[derive(Debug, Clone)]
 enum PyTestDataSource {
@@ -51,14 +52,17 @@ impl From<PyTestDataSource> for TestDataSource {
    }
 }

-/// pyclass containing the struct used for G1
+/// pyclass containing the struct used for G1, this is mostly a helper class
 #[pyclass]
 #[derive(Debug, Clone)]
 struct PyG1 {
    #[pyo3(get, set)]
+    /// Field Element representing x
    x: PyFelt,
    #[pyo3(get, set)]
+    /// Field Element representing y
    y: PyFelt,
+    /// Field Element representing y
    #[pyo3(get, set)]
    z: PyFelt,
 }
@@ -134,39 +138,59 @@ impl pyo3::ToPyObject for PyG1Affine {
    }
 }

-/// pyclass containing the struct used for run_args
+/// Python class containing the struct used for run_args
+///
+/// Returns
+/// -------
+/// PyRunArgs
+///
 #[pyclass]
 #[derive(Clone)]
 struct PyRunArgs {
    #[pyo3(get, set)]
+    /// float: The tolerance for error on model outputs
    pub tolerance: f32,
    #[pyo3(get, set)]
+    /// int: The denominator in the fixed point representation used when quantizing inputs
    pub input_scale: crate::Scale,
    #[pyo3(get, set)]
+    /// int:  The denominator in the fixed point representation used when quantizing parameters
    pub param_scale: crate::Scale,
    #[pyo3(get, set)]
+    /// int: If the scale is ever > scale_rebase_multiplier * input_scale then the scale is rebased to input_scale (this a more advanced parameter, use with caution)
    pub scale_rebase_multiplier: u32,
    #[pyo3(get, set)]
+    /// list[int]: The min and max elements in the lookup table input column
    pub lookup_range: crate::circuit::table::Range,
    #[pyo3(get, set)]
+    /// int: The log_2 number of rows
    pub logrows: u32,
    #[pyo3(get, set)]
+    /// int: The number of inner columns used for the lookup table
    pub num_inner_cols: usize,
    #[pyo3(get, set)]
+    /// string: accepts `public`, `private`, `fixed`, `hashed/public`, `hashed/private`, `polycommit`
    pub input_visibility: Visibility,
    #[pyo3(get, set)]
+    /// string: accepts `public`, `private`, `fixed`, `hashed/public`, `hashed/private`, `polycommit`
    pub output_visibility: Visibility,
    #[pyo3(get, set)]
+    /// string: accepts `public`, `private`, `fixed`, `hashed/public`, `hashed/private`, `polycommit`
    pub param_visibility: Visibility,
    #[pyo3(get, set)]
+    /// list[tuple[str, int]]: Hand-written parser for graph variables, eg. batch_size=1
    pub variables: Vec<(String, usize)>,
    #[pyo3(get, set)]
+    /// bool: Rebase the scale using lookup table for division instead of using a range check
    pub div_rebasing: bool,
    #[pyo3(get, set)]
+    /// bool: Should constants with 0.0 fraction be rebased to scale 0
    pub rebase_frac_zero_constants: bool,
    #[pyo3(get, set)]
+    /// str: check mode, accepts `safe`, `unsafe`
    pub check_mode: CheckMode,
    #[pyo3(get, set)]
+    /// str: commitment type, accepts `kzg`, `ipa`
    pub commitment: PyCommitments,
 }

@@ -197,7 +221,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()),
        }
    }
 }
@@ -226,7 +250,7 @@ impl Into<PyRunArgs> for RunArgs {

 #[pyclass]
 #[derive(Debug, Clone)]
-/// Pyclass marking the type of commitment
+/// pyclass representing an enum, denoting the type of commitment
 pub enum PyCommitments {
    /// KZG commitment
    KZG,
@@ -234,6 +258,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 {
@@ -263,7 +297,19 @@ impl FromStr for PyCommitments {
    }
 }

-/// Converts a felt to big endian
+/// Converts a field element hex string to big endian
+///
+/// Arguments
+/// -------
+/// felt: str
+///     The field element represented as a string
+///
+///
+/// Returns
+/// -------
+/// str
+///     field element represented as a string
+///
 #[pyfunction(signature = (
    felt,
 ))]
@@ -273,22 +319,45 @@ fn felt_to_big_endian(felt: PyFelt) -> PyResult<String> {
 }

 /// Converts a field element hex string to an integer
+///
+/// Arguments
+/// -------
+/// felt: str
+///     The field element represented as a string
+///
+/// Returns
+/// -------
+/// int
+///
 #[pyfunction(signature = (
-    array,
+    felt,
 ))]
-fn felt_to_int(array: PyFelt) -> PyResult<i128> {
-    let felt = crate::pfsys::string_to_field::<Fr>(&array);
+fn felt_to_int(felt: PyFelt) -> PyResult<i128> {
+    let felt = crate::pfsys::string_to_field::<Fr>(&felt);
    let int_rep = felt_to_i128(felt);
    Ok(int_rep)
 }

-/// Converts a field eleement hex string to a floating point number
+/// Converts a field element hex string to a floating point number
+///
+/// Arguments
+/// -------
+/// felt: str
+///    The field element represented as a string
+///
+/// scale: float
+///     The scaling factor used to convert the field element into a floating point representation
+///
+/// Returns
+/// -------
+/// float
+///
 #[pyfunction(signature = (
-    array,
+    felt,
    scale
 ))]
-fn felt_to_float(array: PyFelt, scale: crate::Scale) -> PyResult<f64> {
-    let felt = crate::pfsys::string_to_field::<Fr>(&array);
+fn felt_to_float(felt: PyFelt, scale: crate::Scale) -> PyResult<f64> {
+    let felt = crate::pfsys::string_to_field::<Fr>(&felt);
    let int_rep = felt_to_i128(felt);
    let multiplier = scale_to_multiplier(scale);
    let float_rep = int_rep as f64 / multiplier;
@@ -296,9 +365,23 @@ fn felt_to_float(array: PyFelt, scale: crate::Scale) -> PyResult<f64> {
 }

 /// Converts a floating point element to a field element hex string
+///
+/// Arguments
+/// -------
+/// input: float
+///    The field element represented as a string
+///
+/// scale: float
+///     The scaling factor used to quantize the float into a field element
+///
+/// Returns
+/// -------
+/// str
+///     The field element represented as a string
+///
 #[pyfunction(signature = (
-input,
-scale
+    input,
+    scale
 ))]
 fn float_to_felt(input: f64, scale: crate::Scale) -> PyResult<PyFelt> {
    let int_rep = quantize_float(&input, 0.0, scale)
@@ -308,9 +391,20 @@ fn float_to_felt(input: f64, scale: crate::Scale) -> PyResult<PyFelt> {
 }

 /// Converts a buffer to vector of field elements
+///
+/// Arguments
+/// -------
+/// buffer: list[int]
+///     List of integers representing a buffer
+///
+/// Returns
+/// -------
+/// list[str]
+///     List of field elements represented as strings
+///
 #[pyfunction(signature = (
    buffer
-    ))]
+))]
 fn buffer_to_felts(buffer: Vec<u8>) -> PyResult<Vec<String>> {
    fn u8_array_to_u128_le(arr: [u8; 16]) -> u128 {
        let mut n: u128 = 0;
@@ -369,9 +463,20 @@ fn buffer_to_felts(buffer: Vec<u8>) -> PyResult<Vec<String>> {
 }

 /// Generate a poseidon hash.
+///
+/// Arguments
+/// -------
+/// message: list[str]
+///     List of field elements represnted as strings
+///
+/// Returns
+/// -------
+/// list[str]
+///     List of field elements represented as strings
+///
 #[pyfunction(signature = (
    message,
-    ))]
+))]
 fn poseidon_hash(message: Vec<PyFelt>) -> PyResult<Vec<PyFelt>> {
    let message: Vec<Fr> = message
        .iter()
@@ -392,12 +497,31 @@ fn poseidon_hash(message: Vec<PyFelt>) -> PyResult<Vec<PyFelt>> {
 }

 /// Generate a kzg commitment.
+///
+/// Arguments
+/// -------
+/// message: list[str]
+///     List of field elements represnted as strings
+///
+/// vk_path: str
+///     Path to the verification key
+///
+/// settings_path: str
+///     Path to the settings file
+///
+/// srs_path: str
+///     Path to the Structure Reference String (SRS) file
+///
+/// Returns
+/// -------
+/// list[PyG1Affine]
+///
 #[pyfunction(signature = (
    message,
    vk_path=PathBuf::from(DEFAULT_VK),
    settings_path=PathBuf::from(DEFAULT_SETTINGS),
    srs_path=None
-    ))]
+))]
 fn kzg_commit(
    message: Vec<PyFelt>,
    vk_path: PathBuf,
@@ -432,12 +556,31 @@ fn kzg_commit(
 }

 /// Generate an ipa commitment.
+///
+/// Arguments
+/// -------
+/// message: list[str]
+///     List of field elements represnted as strings
+///
+/// vk_path: str
+///     Path to the verification key
+///
+/// settings_path: str
+///     Path to the settings file
+///
+/// srs_path: str
+///     Path to the Structure Reference String (SRS) file
+///
+/// Returns
+/// -------
+/// list[PyG1Affine]
+///
 #[pyfunction(signature = (
    message,
    vk_path=PathBuf::from(DEFAULT_VK),
    settings_path=PathBuf::from(DEFAULT_SETTINGS),
    srs_path=None
-    ))]
+))]
 fn ipa_commit(
    message: Vec<PyFelt>,
    vk_path: PathBuf,
@@ -472,10 +615,19 @@ fn ipa_commit(
 }

 /// Swap the commitments in a proof
+///
+/// Arguments
+/// -------
+/// proof_path: str
+///     Path to the proof file
+///
+/// witness_path: str
+///     Path to the witness file
+///
 #[pyfunction(signature = (
    proof_path=PathBuf::from(DEFAULT_PROOF),
    witness_path=PathBuf::from(DEFAULT_WITNESS),
-    ))]
+))]
 fn swap_proof_commitments(proof_path: PathBuf, witness_path: PathBuf) -> PyResult<()> {
    crate::execute::swap_proof_commitments_cmd(proof_path, witness_path)
        .map_err(|_| PyIOError::new_err("Failed to swap commitments"))?;
@@ -484,11 +636,27 @@ fn swap_proof_commitments(proof_path: PathBuf, witness_path: PathBuf) -> PyResul
 }

 /// Generates a vk from a pk for a model circuit and saves it to a file
+///
+/// Arguments
+/// -------
+/// path_to_pk: str
+///     Path to the proving key
+///
+/// circuit_settings_path: str
+///     Path to the witness file
+///
+/// vk_output_path: str
+///     Path to create the vk file
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    path_to_pk=PathBuf::from(DEFAULT_PK),
    circuit_settings_path=PathBuf::from(DEFAULT_SETTINGS),
    vk_output_path=PathBuf::from(DEFAULT_VK),
-    ))]
+))]
 fn gen_vk_from_pk_single(
    path_to_pk: PathBuf,
    circuit_settings_path: PathBuf,
@@ -510,10 +678,22 @@ fn gen_vk_from_pk_single(
 }

 /// Generates a vk from a pk for an aggregate circuit and saves it to a file
+///
+/// Arguments
+/// -------
+/// path_to_pk: str
+///     Path to the proving key
+///
+/// vk_output_path: str
+///     Path to create the vk file
+///
+/// Returns
+/// -------
+/// bool
 #[pyfunction(signature = (
    path_to_pk=PathBuf::from(DEFAULT_PK_AGGREGATED),
    vk_output_path=PathBuf::from(DEFAULT_VK_AGGREGATED),
-    ))]
+))]
 fn gen_vk_from_pk_aggr(path_to_pk: PathBuf, vk_output_path: PathBuf) -> PyResult<bool> {
    let pk = load_pk::<KZGCommitmentScheme<Bn256>, AggregationCircuit>(path_to_pk, ())
        .map_err(|_| PyIOError::new_err("Failed to load pk"))?;
@@ -528,6 +708,17 @@ fn gen_vk_from_pk_aggr(path_to_pk: PathBuf, vk_output_path: PathBuf) -> PyResult
 }

 /// Displays the table as a string in python
+///
+/// Arguments
+/// ---------
+/// model: str
+///     Path to the onnx file
+///
+/// Returns
+/// ---------
+/// str
+///     Table of the nodes in the onnx file
+///
 #[pyfunction(signature = (
    model = PathBuf::from(DEFAULT_MODEL),
    py_run_args = None
@@ -543,7 +734,16 @@ fn table(model: PathBuf, py_run_args: Option<PyRunArgs>) -> PyResult<String> {
    }
 }

-/// generates the srs
+/// Generates the Structured Reference String (SRS), use this only for testing purposes
+///
+/// Arguments
+/// ---------
+/// srs_path: str
+///     Path to the create the SRS file
+///
+/// logrows: int
+///     The number of logrows for the SRS file
+///
 #[pyfunction(signature = (
    srs_path,
    logrows,
@@ -554,7 +754,26 @@ fn gen_srs(srs_path: PathBuf, logrows: usize) -> PyResult<()> {
    Ok(())
 }

-/// gets a public srs
+/// Gets a public srs
+///
+/// Arguments
+/// ---------
+/// settings_path: str
+///     Path to the settings file
+///
+/// logrows: int
+///     The number of logrows for the SRS file
+///
+/// srs_path: str
+///     Path to the create the SRS file
+///
+/// commitment: str
+///     Specify the commitment used ("kzg", "ipa")
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    settings_path=PathBuf::from(DEFAULT_SETTINGS),
    logrows=None,
@@ -587,7 +806,23 @@ fn get_srs(
    Ok(true)
 }

-/// generates the circuit settings
+/// Generates the circuit settings
+///
+/// Arguments
+/// ---------
+/// model: str
+///     Path to the onnx file
+///
+/// output: str
+///     Path to create the settings file
+///
+/// py_run_args: PyRunArgs
+///     PyRunArgs object to initialize the settings
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    model=PathBuf::from(DEFAULT_MODEL),
    output=PathBuf::from(DEFAULT_SETTINGS),
@@ -608,7 +843,38 @@ fn gen_settings(
    Ok(true)
 }

-/// calibrates the circuit settings
+/// Calibrates the circuit settings
+///
+/// Arguments
+/// ---------
+/// data: str
+///     Path to the calibration data
+///
+/// model: str
+///     Path to the onnx file
+///
+/// settings: str
+///     Path to the settings file
+///
+/// lookup_safety_margin: int
+///      the lookup safety margin to use for calibration. if the max lookup is 2^k, then the max lookup will be 2^k * lookup_safety_margin. larger = safer but slower
+///
+/// scales: list[int]
+///     Optional scales to specifically try for calibration
+///
+/// scale_rebase_multiplier: list[int]
+///     Optional scale rebase multipliers to specifically try for calibration. This is the multiplier at which we divide to return to the input scale.
+///
+/// max_logrows: int
+///     Optional max logrows to use for calibration
+///
+/// only_range_check_rebase: bool
+///     Check ranges when rebasing
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    data = PathBuf::from(DEFAULT_CALIBRATION_FILE),
    model = PathBuf::from(DEFAULT_MODEL),
@@ -650,7 +916,30 @@ fn calibrate_settings(
    Ok(true)
 }

-/// runs the forward pass operation
+/// Runs the forward pass operation to generate a witness
+///
+/// Arguments
+/// ---------
+/// data: str
+///     Path to the data file
+///
+/// model: str
+///     Path to the compiled model file
+///
+/// output: str
+///     Path to create the witness file
+///
+/// vk_path: str
+///     Path to the verification key
+///
+/// srs_path: str
+///     Path to the SRS file
+///
+/// Returns
+/// -------
+/// dict
+///     Python object containing the witness values
+///
 #[pyfunction(signature = (
    data=PathBuf::from(DEFAULT_DATA),
    model=PathBuf::from(DEFAULT_MODEL),
@@ -677,7 +966,20 @@ fn gen_witness(
    Python::with_gil(|py| Ok(output.to_object(py)))
 }

-/// mocks the prover
+/// Mocks the prover
+///
+/// Arguments
+/// ---------
+/// witness: str
+///     Path to the witness file
+///
+/// model: str
+///     Path to the compiled model file
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    witness=PathBuf::from(DEFAULT_WITNESS),
    model=PathBuf::from(DEFAULT_MODEL),
@@ -690,7 +992,23 @@ fn mock(witness: PathBuf, model: PathBuf) -> PyResult<bool> {
    Ok(true)
 }

-/// mocks the aggregate prover
+/// Mocks the aggregate prover
+///
+/// Arguments
+/// ---------
+/// aggregation_snarks: list[str]
+///     List of paths to the relevant proof files
+///
+/// logrows: int
+///     Number of logrows to use for the aggregation circuit
+///
+/// split_proofs: bool
+///     Indicates whether the accumulated are segments of a larger proof
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    aggregation_snarks=vec![PathBuf::from(DEFAULT_PROOF)],
    logrows=DEFAULT_AGGREGATED_LOGROWS.parse().unwrap(),
@@ -709,7 +1027,32 @@ fn mock_aggregate(
    Ok(true)
 }

-/// runs the prover on a set of inputs
+/// Runs the setup process
+///
+/// Arguments
+/// ---------
+/// model: str
+///     Path to the compiled model file
+///
+/// vk_path: str
+///     Path to create the verification key file
+///
+/// pk_path: str
+///     Path to create the proving key file
+///
+/// srs_path: str
+///     Path to the SRS file
+///
+/// witness_path: str
+///     Path to the witness file
+///
+/// disable_selector_compression: bool
+///     Whether to compress the selectors or not
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    model=PathBuf::from(DEFAULT_MODEL),
    vk_path=PathBuf::from(DEFAULT_VK),
@@ -742,7 +1085,32 @@ fn setup(
    Ok(true)
 }

-/// runs the prover on a set of inputs
+/// Runs the prover on a set of inputs
+///
+/// Arguments
+/// ---------
+/// witness: str
+///     Path to the witness file
+///
+/// model: str
+///     Path to the compiled model file
+///
+/// pk_path: str
+///     Path to the proving key file
+///
+/// proof_path: str
+///     Path to create the proof file
+///
+/// proof_type: str
+///     Accepts `single`, `for-aggr`
+///
+/// srs_path: str
+///     Path to the SRS file
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    witness=PathBuf::from(DEFAULT_WITNESS),
    model=PathBuf::from(DEFAULT_MODEL),
@@ -776,7 +1144,29 @@ fn prove(
    Python::with_gil(|py| Ok(snark.to_object(py)))
 }

-/// verifies a given proof
+/// Verifies a given proof
+///
+/// Arguments
+/// ---------
+/// proof_path: str
+///     Path to create the proof file
+///
+/// settings_path: str
+///     Path to the settings file
+///
+/// vk_path: str
+///     Path to the verification key file
+///
+/// srs_path: str
+///     Path to the SRS file
+///
+/// non_reduced_srs: bool
+///     Whether to reduce the number of SRS logrows to the number of instances rather than the number of logrows used for proofs (only works if the srs were generated in the same ceremony)
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    proof_path=PathBuf::from(DEFAULT_PROOF),
    settings_path=PathBuf::from(DEFAULT_SETTINGS),
@@ -806,6 +1196,38 @@ fn verify(
    Ok(true)
 }

+///  Runs the setup process for an aggregate setup
+///
+/// Arguments
+/// ---------
+/// sample_snarks: list[str]
+///     List of paths to the various proofs
+///
+/// vk_path: str
+///     Path to create the aggregated VK
+///
+/// pk_path: str
+///     Path to create the aggregated PK
+///
+/// logrows: int
+///     Number of logrows to use
+///
+/// split_proofs: bool
+///     Whether the accumulated are segments of a larger proof
+///
+/// srs_path: str
+///     Path to the SRS file
+///
+/// disable_selector_compression: bool
+///     Whether to compress selectors
+///
+/// commitment: str
+///     Accepts `kzg`, `ipa`
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    sample_snarks=vec![PathBuf::from(DEFAULT_PROOF)],
    vk_path=PathBuf::from(DEFAULT_VK_AGGREGATED),
@@ -844,6 +1266,23 @@ fn setup_aggregate(
    Ok(true)
 }

+/// Compiles the circuit for use in other steps
+///
+/// Arguments
+/// ---------
+/// model: str
+///     Path to the onnx model file
+///
+/// compiled_circuit: str
+///     Path to output the compiled circuit
+///
+/// settings_path: str
+///     Path to the settings files
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    model=PathBuf::from(DEFAULT_MODEL),
    compiled_circuit=PathBuf::from(DEFAULT_COMPILED_CIRCUIT),
@@ -862,7 +1301,41 @@ fn compile_circuit(
    Ok(true)
 }

-/// creates an aggregated proof
+/// Creates an aggregated proof
+///
+/// Arguments
+/// ---------
+/// aggregation_snarks: list[str]
+///     List of paths to the various proofs
+///
+/// proof_path: str
+///     Path to output the aggregated proof
+///
+/// vk_path: str
+///     Path to the VK file
+///
+/// transcript:
+///     Proof transcript type to be used. `evm` used by default. `poseidon` is also supported
+///
+/// logrows:
+///     Logrows used for aggregation circuit
+///
+/// check_mode: str
+///     Run sanity checks during calculations. Accepts `safe` or `unsafe`
+///
+/// split-proofs: bool
+///      Whether the accumulated proofs are segments of a larger circuit
+///
+/// srs_path: str
+///     Path to the SRS used
+///
+/// commitment: str
+///     Accepts "kzg" or "ipa"
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    aggregation_snarks=vec![PathBuf::from(DEFAULT_PROOF)],
    proof_path=PathBuf::from(DEFAULT_PROOF_AGGREGATED),
@@ -905,7 +1378,32 @@ fn aggregate(
    Ok(true)
 }

-/// verifies and aggregate proof
+/// Verifies and aggregate proof
+///
+/// Arguments
+/// ---------
+/// proof_path: str
+///      The path to the proof file
+///
+/// vk_path: str
+///     The path to the verification key file
+///
+/// logrows: int
+///     logrows used for aggregation circuit
+///
+/// commitment: str
+///     Accepts "kzg" or "ipa"
+///
+/// reduced_srs: bool
+///     Whether to reduce the number of SRS logrows to the number of instances rather than the number of logrows used for proofs (only works if the srs were generated in the same ceremony)
+///
+/// srs_path: str
+///     The path to the SRS file
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    proof_path=PathBuf::from(DEFAULT_PROOF_AGGREGATED),
    vk_path=PathBuf::from(DEFAULT_VK),
@@ -938,7 +1436,32 @@ fn verify_aggr(
    Ok(true)
 }

-/// creates an EVM compatible verifier, you will need solc installed in your environment to run this
+/// Creates an EVM compatible verifier, you will need solc installed in your environment to run this
+///
+/// Arguments
+/// ---------
+/// vk_path: str
+///     The path to the verification key file
+///
+/// settings_path: str
+///     The path to the settings file
+///
+/// sol_code_path: str
+///     The path to the create the solidity verifier
+///
+/// abi_path: str
+///     The path to create the ABI for the solidity verifier
+///
+/// srs_path: str
+///     The path to the SRS file
+///
+/// render_vk_separately: bool
+///     Whether the verifier key should be rendered as a separate contract. We recommend disabling selector compression if this is enabled. To save the verifier key as a separate contract, set this to true and then call the create-evm-vk command
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    vk_path=PathBuf::from(DEFAULT_VK),
    settings_path=PathBuf::from(DEFAULT_SETTINGS),
@@ -971,7 +1494,26 @@ fn create_evm_verifier(
    Ok(true)
 }

-// creates an EVM compatible data attestation verifier, you will need solc installed in your environment to run this
+/// Creates an EVM compatible data attestation verifier, you will need solc installed in your environment to run this
+///
+/// Arguments
+/// ---------
+/// input_data: str
+///     The path to the .json data file, which should contain the necessary calldata and account addresses needed to read from all the on-chain view functions that return the data that the network ingests as inputs
+///
+/// settings_path: str
+///     The path to the settings file
+///
+/// sol_code_path: str
+///     The path to the create the solidity verifier
+///
+/// abi_path: str
+///     The path to create the ABI for the solidity verifier
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    input_data=PathBuf::from(DEFAULT_DATA),
    settings_path=PathBuf::from(DEFAULT_SETTINGS),
@@ -993,6 +1535,32 @@ fn create_evm_data_attestation(
    Ok(true)
 }

+/// Setup test evm witness
+///
+/// Arguments
+/// ---------
+/// data_path: str
+///     The path to the .json data file, which should include both the network input (possibly private) and the network output (public input to the proof)
+///
+/// compiled_circuit_path: str
+///     The path to the compiled model file (generated using the compile-circuit command)
+///
+/// test_data: str
+///     For testing purposes only. The optional path to the .json data file that will be generated that contains the OnChain data storage information derived from the file information in the data .json file. Should include both the network input (possibly private) and the network output (public input to the proof)
+///
+/// input_sources: str
+///     Where the input data comes from
+///
+/// output_source: str
+///     Where the output data comes from
+///
+/// rpc_url: str
+///     RPC URL for an EVM compatible node, if None, uses Anvil as a local RPC node
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    data_path,
    compiled_circuit_path,
@@ -1027,6 +1595,7 @@ fn setup_test_evm_witness(
    Ok(true)
 }

+/// deploys the solidity verifier
 #[pyfunction(signature = (
    addr_path,
    sol_code_path=PathBuf::from(DEFAULT_SOL_CODE),
@@ -1059,6 +1628,7 @@ fn deploy_evm(
    Ok(true)
 }

+/// deploys the solidity vk verifier
 #[pyfunction(signature = (
    addr_path,
    sol_code_path=PathBuf::from(DEFAULT_VK_SOL),
@@ -1091,6 +1661,7 @@ fn deploy_vk_evm(
    Ok(true)
 }

+/// deploys the solidity da verifier
 #[pyfunction(signature = (
    addr_path,
    input_data,
@@ -1128,6 +1699,27 @@ fn deploy_da_evm(
    Ok(true)
 }
 /// verifies an evm compatible proof, you will need solc installed in your environment to run this
+///
+/// Arguments
+/// ---------
+/// addr_verifier: str
+///     The path to verifier contract's address
+///
+/// proof_path: str
+///     The path to the proof file (generated using the prove command)
+///
+/// rpc_url: str
+///     RPC URL for an Ethereum node, if None will use Anvil but WON'T persist state
+///
+/// addr_da: str
+///     does the verifier use data attestation ?
+///
+/// addr_vk: str
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    addr_verifier,
    proof_path=PathBuf::from(DEFAULT_PROOF),
@@ -1173,7 +1765,35 @@ fn verify_evm(
    Ok(true)
 }

-/// creates an evm compatible aggregate verifier, you will need solc installed in your environment to run this
+/// Creates an evm compatible aggregate verifier, you will need solc installed in your environment to run this
+///
+/// Arguments
+/// ---------
+/// aggregation_settings: str
+///     path to the settings file
+///
+/// vk_path: str
+///     The path to load the desired verification key file
+///
+/// sol_code_path: str
+///     The path to the Solidity code
+///
+/// abi_path: str
+///     The path to output the Solidity verifier ABI
+///
+/// logrows: int
+///     Number of logrows used during aggregated setup
+///
+/// srs_path: str
+///     The path to the SRS file
+///
+/// render_vk_separately: bool
+///     Whether the verifier key should be rendered as a separate contract. We recommend disabling selector compression if this is enabled. To save the verifier key as a separate contract, set this to true and then call the create-evm-vk command
+///
+/// Returns
+/// -------
+/// bool
+///
 #[pyfunction(signature = (
    aggregation_settings=vec![PathBuf::from(DEFAULT_PROOF)],
    vk_path=PathBuf::from(DEFAULT_VK_AGGREGATED),
--- a/src/tensor/mod.rs
+++ b/src/tensor/mod.rs
@@ -934,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]);
@@ -965,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 {
@@ -977,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
@@ -4404,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,3 +1,5 @@
+use core::{iter::FilterMap, slice::Iter};
+
 use crate::circuit::region::ConstantsMap;

 use super::{
@@ -448,25 +450,36 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> 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, .. } => {
-                let map = inner.iter().fold(ConstantsMap::new(), |mut acc, x| {
-                    if let ValType::Constant(c) = x {
-                        acc.insert(*c, x.clone());
+            ValTensor::Value { inner, .. } => inner
+                .par_iter()
+                .filter_map(|x| {
+                    if let ValType::Constant(v) = x {
+                        Some((*v, x.clone()))
+                    } else {
+                        None
                    }
-                    acc
-                });
-                map
-            }
+                })
+                .collect(),
            ValTensor::Instance { .. } => ConstantsMap::new(),
        }
    }
--- 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
github-actions[bot]	1151c15e59	ci: update version string in docs	2024-04-05 17:33:23 +00:00
Jseam	0943e534ee	docs: automated sphinx documentation for python bindings (#714 ) --------- Co-authored-by: dante <45801863+alexander-camuto@users.noreply.github.com> Co-authored-by: Ethan Cemer <tylercemer@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>	2024-04-05 18:33:06 +01:00
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