Repository: EnzymeAD/Enzyme-JAX (language: python)

Issue:
Make JAX import into WORKSPACE work.
The current import was non-functional, and that only remained unnoticed
because there is no BUILD dependency to JAX, so JAX isn't actually
built. (The CI tests install JAX through pip, so the JAX dependency
isn't needed there.)

The fix consists in moving the configuration from JAX's WORKSPACE file
(which is emptied when imported via `http_archive`) to the WORKSPACE
file of this project (plus fixing how its dependency `ducc` is loaded).
Now adding `@jax//jax` to some of the build targets actually (and
successfully) builds JAX.


Discussion / hints:
Something seems stuck. Maybe the build process got stuck, or it might need approval before it starts running?
One of the build targets failed without error (that is visible to me). Can you check what is the problem and restart it?
@wsmoses: Did you retrigger the build? It failed again -- presumable because it never actually ran? If that check isn't supposed to run, maybe you can force-merge the PR? I don't have the permissions...

Fix patch:
diff --git a/.bazelrc b/.bazelrc
index 99943caf..84502e47 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -1,5 +1,6 @@
 build --announce_rc
 
+query --experimental_repo_remote_exec
 build --experimental_repo_remote_exec
 build --cxxopt=-std=c++17 --host_cxxopt=-std=c++17
 build --cxxopt=-w --host_cxxopt=-w
@@ -11,6 +12,7 @@ build --define framework_shared_object=true
 build --define tsl_protobuf_header_only=true
 build --define=use_fast_cpp_protos=true
 build --define=allow_oversize_protos=true
+build --check_visibility=false
 
 build -c opt
 
@@ -27,4 +29,22 @@ build:cuda --@local_config_cuda//:enable_cuda
 build:cuda --@xla//xla/python:jax_cuda_pip_rpaths=true
 # Default hermetic CUDA and CUDNN versions.
 build:cuda --@local_config_cuda//cuda:include_cuda_libs=true
-build:cuda --@local_config_cuda//:cuda_compiler=nvcc
\ No newline at end of file
+build:cuda --@local_config_cuda//:cuda_compiler=nvcc
+
+build:cudaci --repo_env TF_NEED_CUDA=1
+build:cudaci --repo_env TF_NVCC_CLANG=0
+build:cudaci --repo_env CUDA_NVCC=1
+build:cudaci --repo_env CC="/usr/bin/gcc"
+build:cudaci --repo_env TF_NCCL_USE_STUB=1
+build:cudaci --repo_env=HERMETIC_CUDA_VERSION="12.6.1"
+build:cudaci --repo_env=HERMETIC_CUDNN_VERSION="9.4.0"
+# "sm" means we emit only cubin, which is forward compatible within a GPU generation.
+# "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations.
+build:cudaci --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90"
+build:cudaci --crosstool_top="@local_config_cuda//crosstool:toolchain"
+build:cudaci --@local_config_cuda//:enable_cuda
+build:cudaci --@xla//xla/python:jax_cuda_pip_rpaths=true
+# Default hermetic CUDA and CUDNN versions.
+build:cudaci --@local_config_cuda//cuda:include_cuda_libs=true
+build:cudaci --@local_config_cuda//:cuda_compiler=nvcc
+
diff --git a/.buildkite/gpu_pipeline.yml b/.buildkite/gpu_pipeline.yml
index d4d9f84b..5f20577f 100644
--- a/.buildkite/gpu_pipeline.yml
+++ b/.buildkite/gpu_pipeline.yml
@@ -1,37 +1,97 @@
 steps:
   - label: "CUDA"
     agents:
-      queue: "juliagpu"
+      queue: "benchmark"
+      gpu: "rtx4070"
       cuda: "*"
     if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
+    timeout_in_minutes: 180
     commands: |
+      pwd
+      env
       echo "--- Setup :python: Dependencies"
       mkdir -p .local/bin
       export PATH="`pwd`/.local/bin:`pwd`/conda/bin:\$PATH"
       echo "openssl md5 | cut -d' ' -f2" > .local/bin/md5
       chmod +x .local/bin/md5
 
+      # No one tells us what to do
+      unset NV_LIBCUBLAS_VERSION
+      unset NVIDIA_VISIBLE_DEVICES
+      unset NV_NVML_DEV_VERSION
+      unset NV_LIBNCCL_DEV_PACKAGE
+      unset NV_LIBNCCL_DEV_PACKAGE_VERSION
+      unset NVIDIA_REQUIRE_CUDA
+      unset NV_LIBCUBLAS_DEV_PACKAGE
+      unset NV_NVTX_VERSION
+
       curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-amd64
 
       mv bazel* .local/bin/bazel
       chmod +x .local/bin/bazel
+      export PATH="`pwd`/.local/bin:\$PATH"
 
       export RUSTUP_HOME=`pwd`/.local/.rustup
-      export CARGO_HOME=`pwd`/.local/.cargo bash 
+      export CARGO_HOME=`pwd`/.local/.cargo bash
+
       curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
       export PATH="`pwd`/.local/.cargo/bin:\$PATH"
       rustup default stable
       cargo install cxxbridge-cmd
 
       mkdir -p .baztmp
+      touch src/enzyme_ad/jax/deps/tensat/cargo-bazel-lock.json
 
       echo "--- :python: Test"
 
-      # CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --test_output=errors //test/...
-      # CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --cache_test_results=no //test:bench_vs_xla
-      CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --cache_test_results=no --config=cuda //test:llama
+      export CUDA_DIR=`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvcc_cu12/site-packages/nvidia/cuda_nvcc
+      export XLA_FLAGS=--xla_gpu_cuda_data_dir=\$CUDA_DIR
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cusolver_cu12/site-packages/nvidia/cusolver:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cudnn_cu12/site-packages/nvidia/cudnn/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/test.runfiles/pypi_nvidia_cublas_cu12/site-packages/nvidia/cublas/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_cupti_cu12/site-packages/nvidia/cuda_cupti/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_runtime_cu12/site-packages/nvidia/cuda_runtime/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvrtc_cu12/site-packages/nvidia/cuda_nvrtc/lib:\$LD_LIBRARY_PATH"
+      export PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvcc_cu12/site-packages/nvidia/cuda_nvcc/bin:\$PATH"
+      export PATH="`pwd`/bazel-bin/test/llama.runfiles/python_x86_64-unknown-linux-gnu/bin:\$PATH"
+      export TF_CPP_MIN_LOG_LEVEL=0
+
+      BAZEL_BUILD_FLAGS=()
+      BAZEL_BUILD_FLAGS+=(--define=no_aws_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=no_gcp_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=no_hdfs_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=no_kafka_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=no_ignite_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=grpc_no_ares=true)
+
+      BAZEL_BUILD_FLAGS+=(--define=llvm_enable_zlib=false)
+
+      BAZEL_BUILD_FLAGS+=(--verbose_failures)
+      BAZEL_BUILD_FLAGS+=(--cxxopt=-std=c++17 --host_cxxopt=-std=c++17)
+      BAZEL_BUILD_FLAGS+=(--cxxopt=-DTCP_USER_TIMEOUT=0)
+      BAZEL_BUILD_FLAGS+=(--check_visibility=false)
+      BAZEL_BUILD_FLAGS+=(--experimental_cc_shared_library)
+
+      mkdir .eqsat-tmp
+      export TMPDIR="`pwd`/.eqsat-tmp"
+      export TMP=$TMPDIR
+      export TEMP=$TMPDIR
+      BAZEL_BUILD_FLAGS+=(--action_env=TMP=\$TMPDIR --action_env=TEMP=\$TMPDIR --action_env=TMPDIR=\$TMPDIR --sandbox_tmpfs_path=\$TMPDIR)
+      BAZEL_BUILD_FLAGS+=(--action_env=EQSAT_PLATFORM=gpu)
+      BAZEL_BUILD_FLAGS+=(--config=cudaci)
+
+      export EQSAT_PLATFORM=gpu
+
+      CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp run --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL //builddeps:requirements.update || echo "no req update"
+      HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --test_output=errors //test/... || echo "fail1"
+      find `pwd`/bazel-bin/test/llama.runfiles > finds.txt
+      # HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s //test:bench_vs_xla || echo "fail2"
+      HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s \${BAZEL_BUILD_FLAGS[@]} //test:llama || echo "fail3"
+      HERMETIC_PYTHON_VERSION="3.12" bazel-bin/test/llama
       cat bazel-out/*/testlogs/test/llama/test.log
     artifact_paths:
+      - "finds.txt"
+      - "finds2.txt"
       - "bazel-out/*/testlogs/test/llama/test.log"
       - "bazel-out/*/testlogs/test/llama/bench_vs_xla.log"
+      - "bazel-bin/test/llama.runfiles/results_*.csv"
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 0957daca..34f3c73a 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -8,7 +8,7 @@ steps:
         os:
           - macos
         python:
-          - "3.12"
+          - "3.11"
     agents:
       queue: "juliaecosystem"
       os: "{{matrix.os}}"
@@ -55,6 +55,7 @@ steps:
       cargo install cxxbridge-cmd
       
       mkdir -p .baztmp
+      touch src/enzyme_ad/jax/deps/tensat/cargo-bazel-lock.json
       # CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --test_output=errors //test/...
       # CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --cache_test_results=no //test:bench_vs_xla
       CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --cache_test_results=no //test:llama
diff --git a/.gitignore b/.gitignore
index 92eea31c..ade3439c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+*.pyc
 *.swp
 *.swo
 bazel-*
diff --git a/WORKSPACE b/WORKSPACE
index 5a2e6178..7c6745cd 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -26,10 +26,9 @@ http_archive(
 )
 
 load("@jax//third_party/xla:workspace.bzl", "XLA_COMMIT", "XLA_SHA256")
-XLA_COMMIT = "7d4f8d1e8a91e67a713ac69796a22f343d292327"
 http_archive(
     name = "xla",
-    #sha256 = XLA_SHA256,
+    sha256 = XLA_SHA256,
     strip_prefix = "xla-" + XLA_COMMIT,
     urls = ["https://github.com/wsmoses/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)],
     patch_cmds = XLA_PATCHES,
diff --git a/build.sh b/build.sh
index 42f8680f..d7650553 100755
--- a/build.sh
+++ b/build.sh
@@ -23,6 +23,6 @@ export CUDA_HOME=$HOME/miniconda3/
 export PATH=$PATH:$CUDA_HOME/bin
 export CUDACXX=$CUDA_HOME/bin/nvcc
 BAZEL_BUILD_FLAGS+=(--config=cuda)
-HERMETIC_PYTHON_VERSION=3.12 bazel build ${BAZEL_BUILD_FLAGS[@]} :wheel
-pip install bazel-bin/enzyme_ad-0.0.8-py312-none-manylinux2014_x86_64.whl --no-deps --force-reinstall
+HERMETIC_PYTHON_VERSION=3.11 bazel build ${BAZEL_BUILD_FLAGS[@]} :wheel
+pip install bazel-bin/enzyme_ad-0.0.8-py311-none-manylinux2014_x86_64.whl --no-deps --force-reinstall
 
diff --git a/builddeps/BUILD b/builddeps/BUILD
index be0ae35f..451e9e46 100644
--- a/builddeps/BUILD
+++ b/builddeps/BUILD
@@ -10,6 +10,9 @@ compile_pip_requirements(
         "--build-isolation",
         "--rebuild",
     ],
+    extra_deps = [
+        # "@pypi_wheel//:pkg"
+    ],
     requirements_in = "requirements.in",
     requirements_txt = REQUIREMENTS,
     generate_hashes = True,
diff --git a/builddeps/requirements.in b/builddeps/requirements.in
index 58fa9a1e..4a56e860 100644
--- a/builddeps/requirements.in
+++ b/builddeps/requirements.in
@@ -3,6 +3,7 @@
 #
 -r test-requirements.txt
 
-jax >= 0.4.21
-jaxlib >= 0.4.21
+jax
+jaxlib
 absl_py >= 2.0.0
+ortools >= 9.10
diff --git a/builddeps/requirements_lock_3_10.txt b/builddeps/requirements_lock_3_10.txt
index fb0f8bb2..e492db43 100644
--- a/builddeps/requirements_lock_3_10.txt
+++ b/builddeps/requirements_lock_3_10.txt
@@ -10,23 +10,222 @@ absl-py==2.1.0 \
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
-jax==0.4.31 \
+    #   chex
+    #   dm-haiku
+    #   jax-md
+    #   ml-collections
+    #   optax
+    #   orbax-checkpoint
+attrs==24.2.0 \
+    --hash=sha256:5cfb1b9148b5b086569baec03f20d7b6bf3bcacc9a42bebf87ffaaca362f6346 \
+    --hash=sha256:81921eb96de3191c8258c199618104dd27ac608d9366f5e35d011eae1867ede2
+    # via e3nn-jax
+certifi==2024.8.30 \
+    --hash=sha256:922820b53db7a7257ffbda3f597266d435245903d80737e34f8a45ff3e3230d8 \
+    --hash=sha256:bec941d2aa8195e248a60b31ff9f0558284cf01a52591ceda73ea9afffd69fd9
+    # via requests
+charset-normalizer==3.3.2 \
+    --hash=sha256:06435b539f889b1f6f4ac1758871aae42dc3a8c0e24ac9e60c2384973ad73027 \
+    --hash=sha256:06a81e93cd441c56a9b65d8e1d043daeb97a3d0856d177d5c90ba85acb3db087 \
+    --hash=sha256:0a55554a2fa0d408816b3b5cedf0045f4b8e1a6065aec45849de2d6f3f8e9786 \
+    --hash=sha256:0b2b64d2bb6d3fb9112bafa732def486049e63de9618b5843bcdd081d8144cd8 \
+    --hash=sha256:10955842570876604d404661fbccbc9c7e684caf432c09c715ec38fbae45ae09 \
+    --hash=sha256:122c7fa62b130ed55f8f285bfd56d5f4b4a5b503609d181f9ad85e55c89f4185 \
+    --hash=sha256:1ceae2f17a9c33cb48e3263960dc5fc8005351ee19db217e9b1bb15d28c02574 \
+    --hash=sha256:1d3193f4a680c64b4b6a9115943538edb896edc190f0b222e73761716519268e \
+    --hash=sha256:1f79682fbe303db92bc2b1136016a38a42e835d932bab5b3b1bfcfbf0640e519 \
+    --hash=sha256:2127566c664442652f024c837091890cb1942c30937add288223dc895793f898 \
+    --hash=sha256:22afcb9f253dac0696b5a4be4a1c0f8762f8239e21b99680099abd9b2b1b2269 \
+    --hash=sha256:25baf083bf6f6b341f4121c2f3c548875ee6f5339300e08be3f2b2ba1721cdd3 \
+    --hash=sha256:2e81c7b9c8979ce92ed306c249d46894776a909505d8f5a4ba55b14206e3222f \
+    --hash=sha256:3287761bc4ee9e33561a7e058c72ac0938c4f57fe49a09eae428fd88aafe7bb6 \
+    --hash=sha256:34d1c8da1e78d2e001f363791c98a272bb734000fcef47a491c1e3b0505657a8 \
+    --hash=sha256:37e55c8e51c236f95b033f6fb391d7d7970ba5fe7ff453dad675e88cf303377a \
+    --hash=sha256:3d47fa203a7bd9c5b6cee4736ee84ca03b8ef23193c0d1ca99b5089f72645c73 \
+    --hash=sha256:3e4d1f6587322d2788836a99c69062fbb091331ec940e02d12d179c1d53e25fc \
+    --hash=sha256:42cb296636fcc8b0644486d15c12376cb9fa75443e00fb25de0b8602e64c1714 \
+    --hash=sha256:45485e01ff4d3630ec0d9617310448a8702f70e9c01906b0d0118bdf9d124cf2 \
+    --hash=sha256:4a78b2b446bd7c934f5dcedc588903fb2f5eec172f3d29e52a9096a43722adfc \
+    --hash=sha256:4ab2fe47fae9e0f9dee8c04187ce5d09f48eabe611be8259444906793ab7cbce \
+    --hash=sha256:4d0d1650369165a14e14e1e47b372cfcb31d6ab44e6e33cb2d4e57265290044d \
+    --hash=sha256:549a3a73da901d5bc3ce8d24e0600d1fa85524c10287f6004fbab87672bf3e1e \
+    --hash=sha256:55086ee1064215781fff39a1af09518bc9255b50d6333f2e4c74ca09fac6a8f6 \
+    --hash=sha256:572c3763a264ba47b3cf708a44ce965d98555f618ca42c926a9c1616d8f34269 \
+    --hash=sha256:573f6eac48f4769d667c4442081b1794f52919e7edada77495aaed9236d13a96 \
+    --hash=sha256:5b4c145409bef602a690e7cfad0a15a55c13320ff7a3ad7ca59c13bb8ba4d45d \
+    --hash=sha256:6463effa3186ea09411d50efc7d85360b38d5f09b870c48e4600f63af490e56a \
+    --hash=sha256:65f6f63034100ead094b8744b3b97965785388f308a64cf8d7c34f2f2e5be0c4 \
+    --hash=sha256:663946639d296df6a2bb2aa51b60a2454ca1cb29835324c640dafb5ff2131a77 \
+    --hash=sha256:6897af51655e3691ff853668779c7bad41579facacf5fd7253b0133308cf000d \
+    --hash=sha256:68d1f8a9e9e37c1223b656399be5d6b448dea850bed7d0f87a8311f1ff3dabb0 \
+    --hash=sha256:6ac7ffc7ad6d040517be39eb591cac5ff87416c2537df6ba3cba3bae290c0fed \
+    --hash=sha256:6b3251890fff30ee142c44144871185dbe13b11bab478a88887a639655be1068 \
+    --hash=sha256:6c4caeef8fa63d06bd437cd4bdcf3ffefe6738fb1b25951440d80dc7df8c03ac \
+    --hash=sha256:6ef1d82a3af9d3eecdba2321dc1b3c238245d890843e040e41e470ffa64c3e25 \
+    --hash=sha256:753f10e867343b4511128c6ed8c82f7bec3bd026875576dfd88483c5c73b2fd8 \
+    --hash=sha256:7cd13a2e3ddeed6913a65e66e94b51d80a041145a026c27e6bb76c31a853c6ab \
+    --hash=sha256:7ed9e526742851e8d5cc9e6cf41427dfc6068d4f5a3bb03659444b4cabf6bc26 \
+    --hash=sha256:7f04c839ed0b6b98b1a7501a002144b76c18fb1c1850c8b98d458ac269e26ed2 \
+    --hash=sha256:802fe99cca7457642125a8a88a084cef28ff0cf9407060f7b93dca5aa25480db \
+    --hash=sha256:80402cd6ee291dcb72644d6eac93785fe2c8b9cb30893c1af5b8fdd753b9d40f \
+    --hash=sha256:8465322196c8b4d7ab6d1e049e4c5cb460d0394da4a27d23cc242fbf0034b6b5 \
+    --hash=sha256:86216b5cee4b06df986d214f664305142d9c76df9b6512be2738aa72a2048f99 \
+    --hash=sha256:87d1351268731db79e0f8e745d92493ee2841c974128ef629dc518b937d9194c \
+    --hash=sha256:8bdb58ff7ba23002a4c5808d608e4e6c687175724f54a5dade5fa8c67b604e4d \
+    --hash=sha256:8c622a5fe39a48f78944a87d4fb8a53ee07344641b0562c540d840748571b811 \
+    --hash=sha256:8d756e44e94489e49571086ef83b2bb8ce311e730092d2c34ca8f7d925cb20aa \
+    --hash=sha256:8f4a014bc36d3c57402e2977dada34f9c12300af536839dc38c0beab8878f38a \
+    --hash=sha256:9063e24fdb1e498ab71cb7419e24622516c4a04476b17a2dab57e8baa30d6e03 \
+    --hash=sha256:90d558489962fd4918143277a773316e56c72da56ec7aa3dc3dbbe20fdfed15b \
+    --hash=sha256:923c0c831b7cfcb071580d3f46c4baf50f174be571576556269530f4bbd79d04 \
+    --hash=sha256:95f2a5796329323b8f0512e09dbb7a1860c46a39da62ecb2324f116fa8fdc85c \
+    --hash=sha256:96b02a3dc4381e5494fad39be677abcb5e6634bf7b4fa83a6dd3112607547001 \
+    --hash=sha256:9f96df6923e21816da7e0ad3fd47dd8f94b2a5ce594e00677c0013018b813458 \
+    --hash=sha256:a10af20b82360ab00827f916a6058451b723b4e65030c5a18577c8b2de5b3389 \
+    --hash=sha256:a50aebfa173e157099939b17f18600f72f84eed3049e743b68ad15bd69b6bf99 \
+    --hash=sha256:a981a536974bbc7a512cf44ed14938cf01030a99e9b3a06dd59578882f06f985 \
+    --hash=sha256:a9a8e9031d613fd2009c182b69c7b2c1ef8239a0efb1df3f7c8da66d5dd3d537 \
+    --hash=sha256:ae5f4161f18c61806f411a13b0310bea87f987c7d2ecdbdaad0e94eb2e404238 \
+    --hash=sha256:aed38f6e4fb3f5d6bf81bfa990a07806be9d83cf7bacef998ab1a9bd660a581f \
+    --hash=sha256:b01b88d45a6fcb69667cd6d2f7a9aeb4bf53760d7fc536bf679ec94fe9f3ff3d \
+    --hash=sha256:b261ccdec7821281dade748d088bb6e9b69e6d15b30652b74cbbac25e280b796 \
+    --hash=sha256:b2b0a0c0517616b6869869f8c581d4eb2dd83a4d79e0ebcb7d373ef9956aeb0a \
+    --hash=sha256:b4a23f61ce87adf89be746c8a8974fe1c823c891d8f86eb218bb957c924bb143 \
+    --hash=sha256:bd8f7df7d12c2db9fab40bdd87a7c09b1530128315d047a086fa3ae3435cb3a8 \
+    --hash=sha256:beb58fe5cdb101e3a055192ac291b7a21e3b7ef4f67fa1d74e331a7f2124341c \
+    --hash=sha256:c002b4ffc0be611f0d9da932eb0f704fe2602a9a949d1f738e4c34c75b0863d5 \
+    --hash=sha256:c083af607d2515612056a31f0a8d9e0fcb5876b7bfc0abad3ecd275bc4ebc2d5 \
+    --hash=sha256:c180f51afb394e165eafe4ac2936a14bee3eb10debc9d9e4db8958fe36afe711 \
+    --hash=sha256:c235ebd9baae02f1b77bcea61bce332cb4331dc3617d254df3323aa01ab47bd4 \
+    --hash=sha256:cd70574b12bb8a4d2aaa0094515df2463cb429d8536cfb6c7ce983246983e5a6 \
+    --hash=sha256:d0eccceffcb53201b5bfebb52600a5fb483a20b61da9dbc885f8b103cbe7598c \
+    --hash=sha256:d965bba47ddeec8cd560687584e88cf699fd28f192ceb452d1d7ee807c5597b7 \
+    --hash=sha256:db364eca23f876da6f9e16c9da0df51aa4f104a972735574842618b8c6d999d4 \
+    --hash=sha256:ddbb2551d7e0102e7252db79ba445cdab71b26640817ab1e3e3648dad515003b \
+    --hash=sha256:deb6be0ac38ece9ba87dea880e438f25ca3eddfac8b002a2ec3d9183a454e8ae \
+    --hash=sha256:e06ed3eb3218bc64786f7db41917d4e686cc4856944f53d5bdf83a6884432e12 \
+    --hash=sha256:e27ad930a842b4c5eb8ac0016b0a54f5aebbe679340c26101df33424142c143c \
+    --hash=sha256:e537484df0d8f426ce2afb2d0f8e1c3d0b114b83f8850e5f2fbea0e797bd82ae \
+    --hash=sha256:eb00ed941194665c332bf8e078baf037d6c35d7c4f3102ea2d4f16ca94a26dc8 \
+    --hash=sha256:eb6904c354526e758fda7167b33005998fb68c46fbc10e013ca97f21ca5c8887 \
+    --hash=sha256:eb8821e09e916165e160797a6c17edda0679379a4be5c716c260e836e122f54b \
+    --hash=sha256:efcb3f6676480691518c177e3b465bcddf57cea040302f9f4e6e191af91174d4 \
+    --hash=sha256:f27273b60488abe721a075bcca6d7f3964f9f6f067c8c4c605743023d7d3944f \
+    --hash=sha256:f30c3cb33b24454a82faecaf01b19c18562b1e89558fb6c56de4d9118a032fd5 \
+    --hash=sha256:fb69256e180cb6c8a894fee62b3afebae785babc1ee98b81cdf68bbca1987f33 \
+    --hash=sha256:fd1abc0d89e30cc4e02e4064dc67fcc51bd941eb395c502aac3ec19fab46b519 \
+    --hash=sha256:ff8fa367d09b717b2a17a052544193ad76cd49979c805768879cb63d9ca50561
+    # via requests
+chex==0.1.87 \
+    --hash=sha256:0096d89cc8d898bb521ef4bfbf5c24549022b0e5b301f529ab57238896fe6c5d \
+    --hash=sha256:ce536475661fd96d21be0c1728ecdbedd03f8ff950c662dfc338c92ea782cb16
+    # via optax
+contextlib2==21.6.0 \
+    --hash=sha256:3fbdb64466afd23abaf6c977627b75b6139a5a3e8ce38405c5b413aed7a0471f \
+    --hash=sha256:ab1e2bfe1d01d968e1b7e8d9023bc51ef3509bba217bb730cee3827e1ee82869
+    # via ml-collections
+dm-haiku==0.0.12 \
+    --hash=sha256:7448a43a6486bff95253f84e18eacc607d9c1256592573117a9d1d23e2780706 \
+    --hash=sha256:ba0b3acf71433156737fe342c486da11727e5e6c9e054245f4f9b8f0b53eb608
+    # via jax-md
+e3nn-jax==0.20.7 \
+    --hash=sha256:7764edb99d19465e574a23857e5f97420b4d720f06ba8bc18a01f9f2b8301d1c \
+    --hash=sha256:8a1b6f0c00feed565358fceec4c916a623412b6328945fce272f0c33e52b4bfc
+    # via jax-md
+einops==0.8.0 \
+    --hash=sha256:63486517fed345712a8385c100cb279108d9d47e6ae59099b07657e983deae85 \
+    --hash=sha256:9572fb63046264a862693b0a87088af3bdc8c068fde03de63453cbbde245465f
+    # via jax-md
+etils[epath,epy]==1.9.4 \
+    --hash=sha256:4387e7a4911a3b5cc4b92b99a9211386d176b43bae1dac8e2fe345fc2cb95e4b \
+    --hash=sha256:fad950414f0a1ca58c70c70915b0014f9953dd9bcf8aa951a0f75ff9becbeb24
+    # via
+    #   optax
+    #   orbax-checkpoint
+flax==0.9.0 \
+    --hash=sha256:12cd8f7162165ddd56877fb1cd9a4fcb47a31569e4c5343eeb59a36369fa2cfe \
+    --hash=sha256:8b7f361eed0f5324e81f9dc8d02ea53da5f993d7c2e37e7aa5b37d3f6331dd53
+    # via
+    #   dm-haiku
+    #   jax-md
+frozendict==2.4.4 \
+    --hash=sha256:07c3a5dee8bbb84cba770e273cdbf2c87c8e035903af8f781292d72583416801 \
+    --hash=sha256:12a342e439aef28ccec533f0253ea53d75fe9102bd6ea928ff530e76eac38906 \
+    --hash=sha256:1697793b5f62b416c0fc1d94638ec91ed3aa4ab277f6affa3a95216ecb3af170 \
+    --hash=sha256:199a4d32194f3afed6258de7e317054155bc9519252b568d9cfffde7e4d834e5 \
+    --hash=sha256:259528ba6b56fa051bc996f1c4d8b57e30d6dd3bc2f27441891b04babc4b5e73 \
+    --hash=sha256:2b70b431e3a72d410a2cdf1497b3aba2f553635e0c0f657ce311d841bf8273b6 \
+    --hash=sha256:2bd009cf4fc47972838a91e9b83654dc9a095dc4f2bb3a37c3f3124c8a364543 \
+    --hash=sha256:2d8536e068d6bf281f23fa835ac07747fb0f8851879dd189e9709f9567408b4d \
+    --hash=sha256:3148062675536724502c6344d7c485dd4667fdf7980ca9bd05e338ccc0c4471e \
+    --hash=sha256:3f7c031b26e4ee6a3f786ceb5e3abf1181c4ade92dce1f847da26ea2c96008c7 \
+    --hash=sha256:4297d694eb600efa429769125a6f910ec02b85606f22f178bafbee309e7d3ec7 \
+    --hash=sha256:4a59578d47b3949437519b5c39a016a6116b9e787bb19289e333faae81462e59 \
+    --hash=sha256:4ae8d05c8d0b6134bfb6bfb369d5fa0c4df21eabb5ca7f645af95fdc6689678e \
+    --hash=sha256:5d58d9a8d9e49662c6dafbea5e641f97decdb3d6ccd76e55e79818415362ba25 \
+    --hash=sha256:63aa49f1919af7d45fb8fd5dec4c0859bc09f46880bd6297c79bb2db2969b63d \
+    --hash=sha256:6874fec816b37b6eb5795b00e0574cba261bf59723e2de607a195d5edaff0786 \
+    --hash=sha256:6eb716e6a6d693c03b1d53280a1947716129f5ef9bcdd061db5c17dea44b80fe \
+    --hash=sha256:705efca8d74d3facbb6ace80ab3afdd28eb8a237bfb4063ed89996b024bc443d \
+    --hash=sha256:78c94991944dd33c5376f720228e5b252ee67faf3bac50ef381adc9e51e90d9d \
+    --hash=sha256:7f79c26dff10ce11dad3b3627c89bb2e87b9dd5958c2b24325f16a23019b8b94 \
+    --hash=sha256:7fee9420475bb6ff357000092aa9990c2f6182b2bab15764330f4ad7de2eae49 \
+    --hash=sha256:812ab17522ba13637826e65454115a914c2da538356e85f43ecea069813e4b33 \
+    --hash=sha256:85375ec6e979e6373bffb4f54576a68bf7497c350861d20686ccae38aab69c0a \
+    --hash=sha256:87ebcde21565a14fe039672c25550060d6f6d88cf1f339beac094c3b10004eb0 \
+    --hash=sha256:93a7b19afb429cbf99d56faf436b45ef2fa8fe9aca89c49eb1610c3bd85f1760 \
+    --hash=sha256:b3b967d5065872e27b06f785a80c0ed0a45d1f7c9b85223da05358e734d858ca \
+    --hash=sha256:c6bf9260018d653f3cab9bd147bd8592bf98a5c6e338be0491ced3c196c034a3 \
+    --hash=sha256:c8f92425686323a950337da4b75b4c17a3327b831df8c881df24038d560640d4 \
+    --hash=sha256:d13b4310db337f4d2103867c5a05090b22bc4d50ca842093779ef541ea9c9eea \
+    --hash=sha256:d9647563e76adb05b7cde2172403123380871360a114f546b4ae1704510801e5 \
+    --hash=sha256:dc2228874eacae390e63fd4f2bb513b3144066a977dc192163c9f6c7f6de6474 \
+    --hash=sha256:e1b941132d79ce72d562a13341d38fc217bc1ee24d8c35a20d754e79ff99e038 \
+    --hash=sha256:fefeb700bc7eb8b4c2dc48704e4221860d254c8989fb53488540bc44e44a1ac2
+    # via jax-md
+fsspec==2024.9.0 \
+    --hash=sha256:4b0afb90c2f21832df142f292649035d80b421f60a9e1c027802e5a0da2b04e8 \
+    --hash=sha256:a0947d552d8a6efa72cc2c730b12c41d043509156966cca4fb157b0f2a0c574b
+    # via etils
+humanize==4.11.0 \
+    --hash=sha256:b53caaec8532bcb2fff70c8826f904c35943f8cecaca29d272d9df38092736c0 \
+    --hash=sha256:e66f36020a2d5a974c504bd2555cf770621dbdbb6d82f94a6857c0b1ea2608be
+    # via orbax-checkpoint
+idna==3.10 \
+    --hash=sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9 \
+    --hash=sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3
+    # via requests
+importlib-resources==6.4.5 \
+    --hash=sha256:980862a1d16c9e147a59603677fa2aa5fd82b87f223b6cb870695bcfce830065 \
+    --hash=sha256:ac29d5f956f01d5e4bb63102a5a19957f1b9175e45649977264a1416783bb717
+    # via etils
+jax[cuda12-pip]==0.4.31 ; sys_platform == "linux" \
     --hash=sha256:5688703735133d0dc537e99a1d646198a49c9d472d4715fde4bd437c44151bd7 \
     --hash=sha256:fd2d470643a0073d822737f0788f71391656af7e62cc5b2e7995ee390ceac287
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   e3nn-jax
+    #   flax
+    #   jax-md
+    #   jraph
+    #   optax
+    #   orbax-checkpoint
 jax-cuda12-pjrt==0.4.31 \
     --hash=sha256:3e77d1cfebeca06517254eb568f082037e1a2aa3ed8f63c543492ad8ab5a1585 \
     --hash=sha256:8961abb381d893a3c2392ad76ab2067a81f8f2514f3b47d2da3ac24283293fe0
     # via jax-cuda12-plugin
-jax-cuda12-plugin==0.4.31 ; sys_platform == "linux" \
+jax-cuda12-plugin[with-cuda]==0.4.31 \
     --hash=sha256:146f26928ca719a0daa14bb9a9f5a5cbfa20211e76ea05b7bb534277a658a995 \
     --hash=sha256:5048acdf29755303b2887d948137fd82891b48cdbc086e67640ebb976457cfcc \
     --hash=sha256:5cfa46d4106f70c31944c9fafcfa7b04c6f9886a1b35df5477cde7d9c43eb8bd \
     --hash=sha256:7a179e5e80dd9890972d777d597f3c902c6876e42bcf2edcfe4f3ec5a610472e \
     --hash=sha256:a3727a332fbeac625ab6d5ae63d0ed9e62d1e0b5011f72130c9bc96e797395d5 \
     --hash=sha256:cdb6d0c4009438a6a6bd7997ab8a1194beda9ae7322b8d265eea9e551b6c2b4e
+    # via jax
+jax-md @ https://github.com/wsmoses/jax-md/archive/1188490610b95023f8a51166c3f6b92da31e78fe.tar.gz \
+    --hash=sha256:d84b17f0604986364818d5b320e0c6bfa8c898b30095f3f0dc43aeab141d2dde
     # via -r builddeps/test-requirements.txt
 jaxlib==0.4.31 \
     --hash=sha256:185fb615ab6bd95315fbcbd951d84e71f9835d603db8c03c91faee98ce95ff4d \
@@ -47,7 +246,32 @@ jaxlib==0.4.31 \
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   e3nn-jax
     #   jax
+    #   jax-md
+    #   jraph
+    #   optax
+    #   orbax-checkpoint
+jmp==0.0.4 \
+    --hash=sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730 \
+    --hash=sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d
+    # via dm-haiku
+jraph==0.0.6.dev0 \
+    --hash=sha256:350fe37bf717f934f1f84fd3370a480b3178bfcb61dfa217c738971308c57625 \
+    --hash=sha256:c3ac3a0b224b344eb6d367e8bc312d95ea41bf825d01ea31b80dd8c22c0dd8b8
+    # via jax-md
+markdown-it-py==3.0.0 \
+    --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \
+    --hash=sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb
+    # via rich
+mdurl==0.1.2 \
+    --hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
+    --hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
+    # via markdown-it-py
+ml-collections==0.1.1 \
+    --hash=sha256:3fefcc72ec433aa1e5d32307a3e474bbb67f405be814ea52a2166bfc9dbe68cc
+    # via jax-md
 ml-dtypes==0.4.0 \
     --hash=sha256:03e7cda6ef164eed0abb31df69d2c00c3a5ab3e2610b6d4c42183a43329c72a5 \
     --hash=sha256:2bb83fd064db43e67e67d021e547698af4c8d5c6190f2e9b1c53c09f6ff5531d \
@@ -69,6 +293,83 @@ ml-dtypes==0.4.0 \
     # via
     #   jax
     #   jaxlib
+    #   tensorstore
+mpmath==1.3.0 \
+    --hash=sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f \
+    --hash=sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c
+    # via sympy
+msgpack==1.1.0 \
+    --hash=sha256:06f5fd2f6bb2a7914922d935d3b8bb4a7fff3a9a91cfce6d06c13bc42bec975b \
+    --hash=sha256:071603e2f0771c45ad9bc65719291c568d4edf120b44eb36324dcb02a13bfddf \
+    --hash=sha256:0907e1a7119b337971a689153665764adc34e89175f9a34793307d9def08e6ca \
+    --hash=sha256:0f92a83b84e7c0749e3f12821949d79485971f087604178026085f60ce109330 \
+    --hash=sha256:115a7af8ee9e8cddc10f87636767857e7e3717b7a2e97379dc2054712693e90f \
+    --hash=sha256:13599f8829cfbe0158f6456374e9eea9f44eee08076291771d8ae93eda56607f \
+    --hash=sha256:17fb65dd0bec285907f68b15734a993ad3fc94332b5bb21b0435846228de1f39 \
+    --hash=sha256:2137773500afa5494a61b1208619e3871f75f27b03bcfca7b3a7023284140247 \
+    --hash=sha256:3180065ec2abbe13a4ad37688b61b99d7f9e012a535b930e0e683ad6bc30155b \
+    --hash=sha256:398b713459fea610861c8a7b62a6fec1882759f308ae0795b5413ff6a160cf3c \
+    --hash=sha256:3d364a55082fb2a7416f6c63ae383fbd903adb5a6cf78c5b96cc6316dc1cedc7 \
+    --hash=sha256:3df7e6b05571b3814361e8464f9304c42d2196808e0119f55d0d3e62cd5ea044 \
+    --hash=sha256:41c991beebf175faf352fb940bf2af9ad1fb77fd25f38d9142053914947cdbf6 \
+    --hash=sha256:42f754515e0f683f9c79210a5d1cad631ec3d06cea5172214d2176a42e67e19b \
+    --hash=sha256:452aff037287acb1d70a804ffd022b21fa2bb7c46bee884dbc864cc9024128a0 \
+    --hash=sha256:4676e5be1b472909b2ee6356ff425ebedf5142427842aa06b4dfd5117d1ca8a2 \
+    --hash=sha256:46c34e99110762a76e3911fc923222472c9d681f1094096ac4102c18319e6468 \
+    --hash=sha256:471e27a5787a2e3f974ba023f9e265a8c7cfd373632247deb225617e3100a3c7 \
+    --hash=sha256:4a1964df7b81285d00a84da4e70cb1383f2e665e0f1f2a7027e683956d04b734 \
+    --hash=sha256:4b51405e36e075193bc051315dbf29168d6141ae2500ba8cd80a522964e31434 \
+    --hash=sha256:4d1b7ff2d6146e16e8bd665ac726a89c74163ef8cd39fa8c1087d4e52d3a2325 \
+    --hash=sha256:53258eeb7a80fc46f62fd59c876957a2d0e15e6449a9e71842b6d24419d88ca1 \
+    --hash=sha256:534480ee5690ab3cbed89d4c8971a5c631b69a8c0883ecfea96c19118510c846 \
+    --hash=sha256:58638690ebd0a06427c5fe1a227bb6b8b9fdc2bd07701bec13c2335c82131a88 \
+    --hash=sha256:58dfc47f8b102da61e8949708b3eafc3504509a5728f8b4ddef84bd9e16ad420 \
+    --hash=sha256:59caf6a4ed0d164055ccff8fe31eddc0ebc07cf7326a2aaa0dbf7a4001cd823e \
+    --hash=sha256:5dbad74103df937e1325cc4bfeaf57713be0b4f15e1c2da43ccdd836393e2ea2 \
+    --hash=sha256:5e1da8f11a3dd397f0a32c76165cf0c4eb95b31013a94f6ecc0b280c05c91b59 \
+    --hash=sha256:646afc8102935a388ffc3914b336d22d1c2d6209c773f3eb5dd4d6d3b6f8c1cb \
+    --hash=sha256:64fc9068d701233effd61b19efb1485587560b66fe57b3e50d29c5d78e7fef68 \
+    --hash=sha256:65553c9b6da8166e819a6aa90ad15288599b340f91d18f60b2061f402b9a4915 \
+    --hash=sha256:685ec345eefc757a7c8af44a3032734a739f8c45d1b0ac45efc5d8977aa4720f \
+    --hash=sha256:6ad622bf7756d5a497d5b6836e7fc3752e2dd6f4c648e24b1803f6048596f701 \
+    --hash=sha256:73322a6cc57fcee3c0c57c4463d828e9428275fb85a27aa2aa1a92fdc42afd7b \
+    --hash=sha256:74bed8f63f8f14d75eec75cf3d04ad581da6b914001b474a5d3cd3372c8cc27d \
+    --hash=sha256:79ec007767b9b56860e0372085f8504db5d06bd6a327a335449508bbee9648fa \
+    --hash=sha256:7a946a8992941fea80ed4beae6bff74ffd7ee129a90b4dd5cf9c476a30e9708d \
+    --hash=sha256:7ad442d527a7e358a469faf43fda45aaf4ac3249c8310a82f0ccff9164e5dccd \
+    --hash=sha256:7c9a35ce2c2573bada929e0b7b3576de647b0defbd25f5139dcdaba0ae35a4cc \
+    --hash=sha256:7e7b853bbc44fb03fbdba34feb4bd414322180135e2cb5164f20ce1c9795ee48 \
+    --hash=sha256:879a7b7b0ad82481c52d3c7eb99bf6f0645dbdec5134a4bddbd16f3506947feb \
+    --hash=sha256:8a706d1e74dd3dea05cb54580d9bd8b2880e9264856ce5068027eed09680aa74 \
+    --hash=sha256:8a84efb768fb968381e525eeeb3d92857e4985aacc39f3c47ffd00eb4509315b \
+    --hash=sha256:8cf9e8c3a2153934a23ac160cc4cba0ec035f6867c8013cc6077a79823370346 \
+    --hash=sha256:8da4bf6d54ceed70e8861f833f83ce0814a2b72102e890cbdfe4b34764cdd66e \
+    --hash=sha256:8e59bca908d9ca0de3dc8684f21ebf9a690fe47b6be93236eb40b99af28b6ea6 \
+    --hash=sha256:914571a2a5b4e7606997e169f64ce53a8b1e06f2cf2c3a7273aa106236d43dd5 \
+    --hash=sha256:a51abd48c6d8ac89e0cfd4fe177c61481aca2d5e7ba42044fd218cfd8ea9899f \
+    --hash=sha256:a52a1f3a5af7ba1c9ace055b659189f6c669cf3657095b50f9602af3a3ba0fe5 \
+    --hash=sha256:ad33e8400e4ec17ba782f7b9cf868977d867ed784a1f5f2ab46e7ba53b6e1e1b \
+    --hash=sha256:b4c01941fd2ff87c2a934ee6055bda4ed353a7846b8d4f341c428109e9fcde8c \
+    --hash=sha256:bce7d9e614a04d0883af0b3d4d501171fbfca038f12c77fa838d9f198147a23f \
+    --hash=sha256:c40ffa9a15d74e05ba1fe2681ea33b9caffd886675412612d93ab17b58ea2fec \
+    --hash=sha256:c5a91481a3cc573ac8c0d9aace09345d989dc4a0202b7fcb312c88c26d4e71a8 \
+    --hash=sha256:c921af52214dcbb75e6bdf6a661b23c3e6417f00c603dd2070bccb5c3ef499f5 \
+    --hash=sha256:d46cf9e3705ea9485687aa4001a76e44748b609d260af21c4ceea7f2212a501d \
+    --hash=sha256:d8ce0b22b890be5d252de90d0e0d119f363012027cf256185fc3d474c44b1b9e \
+    --hash=sha256:dd432ccc2c72b914e4cb77afce64aab761c1137cc698be3984eee260bcb2896e \
+    --hash=sha256:e0856a2b7e8dcb874be44fea031d22e5b3a19121be92a1e098f46068a11b0870 \
+    --hash=sha256:e1f3c3d21f7cf67bcf2da8e494d30a75e4cf60041d98b3f79875afb5b96f3a3f \
+    --hash=sha256:f1ba6136e650898082d9d5a5217d5906d1e138024f836ff48691784bbe1adf96 \
+    --hash=sha256:f3e9b4936df53b970513eac1758f3882c88658a220b58dcc1e39606dccaaf01c \
+    --hash=sha256:f80bc7d47f76089633763f952e67f8214cb7b3ee6bfa489b3cb6a84cfac114cd \
+    --hash=sha256:fd2906780f25c8ed5d7b323379f6138524ba793428db5d0e9d226d3fa6aa1788
+    # via
+    #   flax
+    #   orbax-checkpoint
+nest-asyncio==1.6.0 \
+    --hash=sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe \
+    --hash=sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c
+    # via orbax-checkpoint
 numpy==2.1.0 \
     --hash=sha256:08801848a40aea24ce16c2ecde3b756f9ad756586fb2d13210939eb69b023f5b \
     --hash=sha256:0937e54c09f7a9a68da6889362ddd2ff584c02d015ec92672c099b61555f8911 \
@@ -124,15 +425,175 @@ numpy==2.1.0 \
     --hash=sha256:ffbd6faeb190aaf2b5e9024bac9622d2ee549b7ec89ef3a9373fa35313d44e0e
     # via
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   dm-haiku
+    #   e3nn-jax
     #   jax
+    #   jax-md
     #   jaxlib
+    #   jmp
+    #   jraph
     #   ml-dtypes
     #   opt-einsum
+    #   optax
+    #   orbax-checkpoint
     #   scipy
+    #   tensorstore
+nvidia-cublas-cu12==12.6.3.3 \
+    --hash=sha256:e1f70bee38b964eac1907293b336bceb24498a4243e61eaf91a52977c59aebc4 \
+    --hash=sha256:e531199ca4f1f764fb45bc1dde49a006f6765033f9c89c737e4553b9502ca1f5 \
+    --hash=sha256:f33fb68e101d99470c82d17f92a0dd9f74de2a21685c217f4716cdd63b1316eb
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+nvidia-cuda-cupti-cu12==12.6.80 \
+    --hash=sha256:358b4a1d35370353d52e12f0a7d1769fc01ff74a191689d3870b2123156184c4 \
+    --hash=sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73 \
+    --hash=sha256:bbe6ae76e83ce5251b56e8c8e61a964f757175682bbad058b170b136266ab00a
+    # via jax-cuda12-plugin
+nvidia-cuda-nvcc-cu12==12.6.77 \
+    --hash=sha256:26e118b4dae4ebd316ab6f5ee3f3c4cb929f8265f717a8e016ba57e254c0b4f1 \
+    --hash=sha256:9a2aad43b8f4e4be9c237ece04f64450df77871eeebffad5a5a872714a97ca77 \
+    --hash=sha256:b66bb5dd6b8ae62262586691977d3b4a425e91db61b4b9bc2f6b42bcd4154b96
+    # via jax-cuda12-plugin
+nvidia-cuda-runtime-cu12==12.6.77 \
+    --hash=sha256:86c58044c824bf3c173c49a2dbc7a6c8b53cb4e4dca50068be0bf64e9dab3f7f \
+    --hash=sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8 \
+    --hash=sha256:d461264ecb429c84c8879a7153499ddc7b19b5f8d84c204307491989a365588e
+    # via jax-cuda12-plugin
+nvidia-cudnn-cu12==9.4.0.58 \
+    --hash=sha256:2b891934e9e9bfcce61ce7f123ac3d9fd3a3d21ee9f46426aa7a96a6e9452410 \
+    --hash=sha256:cbbc57da0cbab1f7f3a9b7790f702b75c9adb00ee67499e84dba2b458065749b \
+    --hash=sha256:eabcbb94ebc0c8eb12e42133e4aaec6b2d73750beb144d360f7f60fec81489e1
+    # via jax-cuda12-plugin
+nvidia-cufft-cu12==11.3.0.4 \
+    --hash=sha256:6048ebddfb90d09d2707efb1fd78d4e3a77cb3ae4dc60e19aab6be0ece2ae464 \
+    --hash=sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca \
+    --hash=sha256:8510990de9f96c803a051822618d42bf6cb8f069ff3f48d93a8486efdacb48fb
+    # via jax-cuda12-plugin
+nvidia-cusolver-cu12==11.7.1.2 \
+    --hash=sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0 \
+    --hash=sha256:6813f9d8073f555444a8705f3ab0296d3e1cb37a16d694c5fc8b862a0d8706d7 \
+    --hash=sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6
+    # via jax-cuda12-plugin
+nvidia-cusparse-cu12==12.5.4.2 \
+    --hash=sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f \
+    --hash=sha256:4acb8c08855a26d737398cba8fb6f8f5045d93f82612b4cfd84645a2332ccf20 \
+    --hash=sha256:7aa32fa5470cf754f72d1116c7cbc300b4e638d3ae5304cfa4a638a5b87161b1
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cusolver-cu12
+nvidia-nccl-cu12==2.23.4 \
+    --hash=sha256:aa946c8327e22ced28e7cef508a334673abc42064ec85f02d005ba1785ea4cec \
+    --hash=sha256:b097258d9aab2fa9f686e33c6fe40ae57b27df60cedbd15d139701bb5509e0c1
+    # via jax-cuda12-plugin
+nvidia-nvjitlink-cu12==12.6.77 \
+    --hash=sha256:3bf10d85bb1801e9c894c6e197e44dd137d2a0a9e43f8450e9ad13f2df0dd52d \
+    --hash=sha256:410718cd44962bed862a31dd0318620f6f9a8b28a6291967bcfcb446a6516771 \
+    --hash=sha256:9ae346d16203ae4ea513be416495167a0101d33d2d14935aa9c1829a3fb45142
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cufft-cu12
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
 opt-einsum==3.3.0 \
     --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
     --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
     # via jax
+optax==0.2.3 \
+    --hash=sha256:083e603dcd731d7e74d99f71c12f77937dd53f79001b4c09c290e4f47dd2e94f \
+    --hash=sha256:ec7ab925440b0c5a512e1f24fba0fb3e7d760a7fd5d2496d7a691e9d37da01d9
+    # via
+    #   flax
+    #   jax-md
+orbax-checkpoint==0.6.4 \
+    --hash=sha256:366b4d528a7322e1b3d9ddcaed45c8515add0d2fc69c8975c30d98638543240f \
+    --hash=sha256:b4f2608ee4d1da67f7619fc35ff9c928ecdf4ccf7546eeb43ecf38c2608b6dea
+    # via flax
+protobuf==5.28.2 \
+    --hash=sha256:2c69461a7fcc8e24be697624c09a839976d82ae75062b11a0972e41fd2cd9132 \
+    --hash=sha256:35cfcb15f213449af7ff6198d6eb5f739c37d7e4f1c09b5d0641babf2cc0c68f \
+    --hash=sha256:52235802093bd8a2811abbe8bf0ab9c5f54cca0a751fdd3f6ac2a21438bffece \
+    --hash=sha256:59379674ff119717404f7454647913787034f03fe7049cbef1d74a97bb4593f0 \
+    --hash=sha256:5e8a95246d581eef20471b5d5ba010d55f66740942b95ba9b872d918c459452f \
+    --hash=sha256:87317e9bcda04a32f2ee82089a204d3a2f0d3c8aeed16568c7daf4756e4f1fe0 \
+    --hash=sha256:8ddc60bf374785fb7cb12510b267f59067fa10087325b8e1855b898a0d81d276 \
+    --hash=sha256:a8b9403fc70764b08d2f593ce44f1d2920c5077bf7d311fefec999f8c40f78b7 \
+    --hash=sha256:c0ea0123dac3399a2eeb1a1443d82b7afc9ff40241433296769f7da42d142ec3 \
+    --hash=sha256:ca53faf29896c526863366a52a8f4d88e69cd04ec9571ed6082fa117fac3ab36 \
+    --hash=sha256:eeea10f3dc0ac7e6b4933d32db20662902b4ab81bf28df12218aa389e9c2102d
+    # via orbax-checkpoint
+pygments==2.18.0 \
+    --hash=sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199 \
+    --hash=sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a
+    # via rich
+pyyaml==6.0.2 \
+    --hash=sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff \
+    --hash=sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48 \
+    --hash=sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086 \
+    --hash=sha256:0b69e4ce7a131fe56b7e4d770c67429700908fc0752af059838b1cfb41960e4e \
+    --hash=sha256:0ffe8360bab4910ef1b9e87fb812d8bc0a308b0d0eef8c8f44e0254ab3b07133 \
+    --hash=sha256:11d8f3dd2b9c1207dcaf2ee0bbbfd5991f571186ec9cc78427ba5bd32afae4b5 \
+    --hash=sha256:17e311b6c678207928d649faa7cb0d7b4c26a0ba73d41e99c4fff6b6c3276484 \
+    --hash=sha256:1e2120ef853f59c7419231f3bf4e7021f1b936f6ebd222406c3b60212205d2ee \
+    --hash=sha256:1f71ea527786de97d1a0cc0eacd1defc0985dcf6b3f17bb77dcfc8c34bec4dc5 \
+    --hash=sha256:23502f431948090f597378482b4812b0caae32c22213aecf3b55325e049a6c68 \
+    --hash=sha256:24471b829b3bf607e04e88d79542a9d48bb037c2267d7927a874e6c205ca7e9a \
+    --hash=sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf \
+    --hash=sha256:2e99c6826ffa974fe6e27cdb5ed0021786b03fc98e5ee3c5bfe1fd5015f42b99 \
+    --hash=sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8 \
+    --hash=sha256:3ad2a3decf9aaba3d29c8f537ac4b243e36bef957511b4766cb0057d32b0be85 \
+    --hash=sha256:3b1fdb9dc17f5a7677423d508ab4f243a726dea51fa5e70992e59a7411c89d19 \
+    --hash=sha256:41e4e3953a79407c794916fa277a82531dd93aad34e29c2a514c2c0c5fe971cc \
+    --hash=sha256:43fa96a3ca0d6b1812e01ced1044a003533c47f6ee8aca31724f78e93ccc089a \
+    --hash=sha256:50187695423ffe49e2deacb8cd10510bc361faac997de9efef88badc3bb9e2d1 \
+    --hash=sha256:5ac9328ec4831237bec75defaf839f7d4564be1e6b25ac710bd1a96321cc8317 \
+    --hash=sha256:5d225db5a45f21e78dd9358e58a98702a0302f2659a3c6cd320564b75b86f47c \
+    --hash=sha256:6395c297d42274772abc367baaa79683958044e5d3835486c16da75d2a694631 \
+    --hash=sha256:688ba32a1cffef67fd2e9398a2efebaea461578b0923624778664cc1c914db5d \
+    --hash=sha256:68ccc6023a3400877818152ad9a1033e3db8625d899c72eacb5a668902e4d652 \
+    --hash=sha256:70b189594dbe54f75ab3a1acec5f1e3faa7e8cf2f1e08d9b561cb41b845f69d5 \
+    --hash=sha256:797b4f722ffa07cc8d62053e4cff1486fa6dc094105d13fea7b1de7d8bf71c9e \
+    --hash=sha256:7c36280e6fb8385e520936c3cb3b8042851904eba0e58d277dca80a5cfed590b \
+    --hash=sha256:7e7401d0de89a9a855c839bc697c079a4af81cf878373abd7dc625847d25cbd8 \
+    --hash=sha256:80bab7bfc629882493af4aa31a4cfa43a4c57c83813253626916b8c7ada83476 \
+    --hash=sha256:82d09873e40955485746739bcb8b4586983670466c23382c19cffecbf1fd8706 \
+    --hash=sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563 \
+    --hash=sha256:8824b5a04a04a047e72eea5cec3bc266db09e35de6bdfe34c9436ac5ee27d237 \
+    --hash=sha256:8b9c7197f7cb2738065c481a0461e50ad02f18c78cd75775628afb4d7137fb3b \
+    --hash=sha256:9056c1ecd25795207ad294bcf39f2db3d845767be0ea6e6a34d856f006006083 \
+    --hash=sha256:936d68689298c36b53b29f23c6dbb74de12b4ac12ca6cfe0e047bedceea56180 \
+    --hash=sha256:9b22676e8097e9e22e36d6b7bda33190d0d400f345f23d4065d48f4ca7ae0425 \
+    --hash=sha256:a4d3091415f010369ae4ed1fc6b79def9416358877534caf6a0fdd2146c87a3e \
+    --hash=sha256:a8786accb172bd8afb8be14490a16625cbc387036876ab6ba70912730faf8e1f \
+    --hash=sha256:a9f8c2e67970f13b16084e04f134610fd1d374bf477b17ec1599185cf611d725 \
+    --hash=sha256:bc2fa7c6b47d6bc618dd7fb02ef6fdedb1090ec036abab80d4681424b84c1183 \
+    --hash=sha256:c70c95198c015b85feafc136515252a261a84561b7b1d51e3384e0655ddf25ab \
+    --hash=sha256:cc1c1159b3d456576af7a3e4d1ba7e6924cb39de8f67111c735f6fc832082774 \
+    --hash=sha256:ce826d6ef20b1bc864f0a68340c8b3287705cae2f8b4b1d932177dcc76721725 \
+    --hash=sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e \
+    --hash=sha256:d7fded462629cfa4b685c5416b949ebad6cec74af5e2d42905d41e257e0869f5 \
+    --hash=sha256:d84a1718ee396f54f3a086ea0a66d8e552b2ab2017ef8b420e92edbc841c352d \
+    --hash=sha256:d8e03406cac8513435335dbab54c0d385e4a49e4945d2909a581c83647ca0290 \
+    --hash=sha256:e10ce637b18caea04431ce14fabcf5c64a1c61ec9c56b071a4b7ca131ca52d44 \
+    --hash=sha256:ec031d5d2feb36d1d1a24380e4db6d43695f3748343d99434e6f5f9156aaa2ed \
+    --hash=sha256:ef6107725bd54b262d6dedcc2af448a266975032bc85ef0172c5f059da6325b4 \
+    --hash=sha256:efdca5630322a10774e8e98e1af481aad470dd62c3170801852d752aa7a783ba \
+    --hash=sha256:f753120cb8181e736c57ef7636e83f31b9c0d1722c516f7e86cf15b7aa57ff12 \
+    --hash=sha256:ff3824dc5261f50c9b0dfb3be22b4567a6f938ccce4587b38952d85fd9e9afe4
+    # via
+    #   flax
+    #   ml-collections
+    #   orbax-checkpoint
+requests==2.32.3 ; sys_platform == "linux" \
+    --hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \
+    --hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6
+    # via -r builddeps/test-requirements.txt
+rich==13.9.2 \
+    --hash=sha256:51a2c62057461aaf7152b4d611168f93a9fc73068f8ded2790f29fe2b5366d0c \
+    --hash=sha256:8c82a3d3f8dcfe9e734771313e606b39d8247bb6b826e196f4914b333b743cf1
+    # via flax
 scipy==1.14.1 \
     --hash=sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e \
     --hash=sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79 \
@@ -170,3 +631,61 @@ scipy==1.14.1 \
     # via
     #   jax
     #   jaxlib
+six==1.16.0 \
+    --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
+    --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
+    # via ml-collections
+sympy==1.13.3 \
+    --hash=sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73 \
+    --hash=sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9
+    # via e3nn-jax
+tabulate==0.9.0 \
+    --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \
+    --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f
+    # via dm-haiku
+tensorstore==0.1.66 \
+    --hash=sha256:0261b1b7121186ac4898c9b4e29289c46754f9d7d73e906620e9a21386fe564e \
+    --hash=sha256:06f7dba630d75f65201b1c021fcd89e5f8538506b444b5dd712f35c11b812fc8 \
+    --hash=sha256:0d1f8c170d11e6e7a2132e3e0bb34565d24faff1fca42b5c7042772356e2fce0 \
+    --hash=sha256:0de07f34b181023f9179b1e2f7cb255dfad300a4d3f18639d5c4cc955ae88b6b \
+    --hash=sha256:16c4474b10c90b27d422ccffd31f78ed7c2f2976330ac80a4b53a73fd8f041d9 \
+    --hash=sha256:356e2a7fa256060c9da5d88cac6d9384d99887ab0df819c61f1dbe298ab56085 \
+    --hash=sha256:3bd37baa5983b9707c2b5e63301a2091aac3c21fef9891b5fce67597f9cfd61e \
+    --hash=sha256:486be8ec698d2eb10f297353351ff69e90551297185a158d6875e097fa08de68 \
+    --hash=sha256:5af8bf532ae4cc75709093483043ccd1d0f480d0f2b5e965ee5724f7d229658c \
+    --hash=sha256:5ca2a389b8b274f3a8bef907469bc07f700001ee8ae0bc63a5164135c5afa8ee \
+    --hash=sha256:6de7521a0f1f0fb4c4bb80ba945cdd926c4301d1ff2d821946f1675c0f532f5c \
+    --hash=sha256:94523e656c013721924cb5a2955b6c80f3bc4b469dd3d89cb759ab38f1c96cac \
+    --hash=sha256:9c7a672579b08a9c6af4875c56ced1f61f19df4cabe7b61b445a07307fe85d0a \
+    --hash=sha256:a150257280f580f73b3c037a8e60c56a51e5c81d4bc3a1a90cab0e00a292dddd \
+    --hash=sha256:a92dfc7b05e959e5823a87b8a20dc291382d5a647dff49eb560d5516e3490346 \
+    --hash=sha256:afcfa90c5dfc969e2286a30df97b91588486f77b232793abdd7c2cfd20df29d5 \
+    --hash=sha256:b77ee47da8a1b3d6fd03e23a8f853a2a666037f03e21546b4eb2b4cf43e13a96 \
+    --hash=sha256:bb3d9bc9156adaf2212953d941efd993b57a8da7053b93f512c1c313dfdddb89 \
+    --hash=sha256:bbc662c4697cd21443518ffb0c8f23aa72073c6a80f9828aea34a3b43d2230dc \
+    --hash=sha256:c4047f37ea0fbf665e25e918ad9d178fb2b3513203654ae10d8853d14e10a107 \
+    --hash=sha256:c974926fb2f3048fa3d36fc2fdf572403c7b70191193e83c5af78fdc46910b62
+    # via
+    #   flax
+    #   orbax-checkpoint
+toolz==1.0.0 \
+    --hash=sha256:292c8f1c4e7516bf9086f8850935c799a874039c8bcf959d47b600e4c44a6236 \
+    --hash=sha256:2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02
+    # via chex
+typing-extensions==4.12.2 \
+    --hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \
+    --hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8
+    # via
+    #   chex
+    #   etils
+    #   flax
+    #   orbax-checkpoint
+    #   rich
+urllib3==2.2.3 \
+    --hash=sha256:ca899ca043dcb1bafa3e262d73aa25c465bfb49e0bd9dd5d59f1d0acba2f8fac \
+    --hash=sha256:e7d814a81dad81e6caf2ec9fdedb284ecc9c73076b62654547cc64ccdcae26e9
+    # via requests
+zipp==3.20.2 \
+    --hash=sha256:a817ac80d6cf4b23bf7f2828b7cabf326f15a001bea8b1f9b49631780ba28350 \
+    --hash=sha256:bc9eb26f4506fda01b81bcde0ca78103b6e62f991b381fec825435c836edbc29
+    # via etils
diff --git a/builddeps/requirements_lock_3_11.txt b/builddeps/requirements_lock_3_11.txt
index b64d3b71..3d15e356 100644
--- a/builddeps/requirements_lock_3_11.txt
+++ b/builddeps/requirements_lock_3_11.txt
@@ -4,14 +4,22 @@
 #
 #    bazel run //builddeps:requirements.update
 #
---find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
-
 absl-py==2.1.0 \
     --hash=sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308 \
     --hash=sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   dm-haiku
+    #   jax-md
+    #   ml-collections
+    #   optax
+    #   orbax-checkpoint
+attrs==24.2.0 \
+    --hash=sha256:5cfb1b9148b5b086569baec03f20d7b6bf3bcacc9a42bebf87ffaaca362f6346 \
+    --hash=sha256:81921eb96de3191c8258c199618104dd27ac608d9366f5e35d011eae1867ede2
+    # via e3nn-jax
 certifi==2024.7.4 \
     --hash=sha256:5a1e7645bc0ec61a09e26c36f6106dd4cf40c6db3a1fb6352b0244e7fb057c7b \
     --hash=sha256:c198e21b1289c2ab85ee4e67bb4b4ef3ead0892059901a8d5b622f24a1101e90
@@ -108,27 +116,116 @@ charset-normalizer==3.3.2 \
     --hash=sha256:fd1abc0d89e30cc4e02e4064dc67fcc51bd941eb395c502aac3ec19fab46b519 \
     --hash=sha256:ff8fa367d09b717b2a17a052544193ad76cd49979c805768879cb63d9ca50561
     # via requests
+chex==0.1.87 \
+    --hash=sha256:0096d89cc8d898bb521ef4bfbf5c24549022b0e5b301f529ab57238896fe6c5d \
+    --hash=sha256:ce536475661fd96d21be0c1728ecdbedd03f8ff950c662dfc338c92ea782cb16
+    # via optax
+contextlib2==21.6.0 \
+    --hash=sha256:3fbdb64466afd23abaf6c977627b75b6139a5a3e8ce38405c5b413aed7a0471f \
+    --hash=sha256:ab1e2bfe1d01d968e1b7e8d9023bc51ef3509bba217bb730cee3827e1ee82869
+    # via ml-collections
+dm-haiku==0.0.12 \
+    --hash=sha256:7448a43a6486bff95253f84e18eacc607d9c1256592573117a9d1d23e2780706 \
+    --hash=sha256:ba0b3acf71433156737fe342c486da11727e5e6c9e054245f4f9b8f0b53eb608
+    # via jax-md
+e3nn-jax==0.20.7 \
+    --hash=sha256:7764edb99d19465e574a23857e5f97420b4d720f06ba8bc18a01f9f2b8301d1c \
+    --hash=sha256:8a1b6f0c00feed565358fceec4c916a623412b6328945fce272f0c33e52b4bfc
+    # via jax-md
+einops==0.8.0 \
+    --hash=sha256:63486517fed345712a8385c100cb279108d9d47e6ae59099b07657e983deae85 \
+    --hash=sha256:9572fb63046264a862693b0a87088af3bdc8c068fde03de63453cbbde245465f
+    # via jax-md
+etils[epath,epy]==1.9.4 \
+    --hash=sha256:4387e7a4911a3b5cc4b92b99a9211386d176b43bae1dac8e2fe345fc2cb95e4b \
+    --hash=sha256:fad950414f0a1ca58c70c70915b0014f9953dd9bcf8aa951a0f75ff9becbeb24
+    # via
+    #   optax
+    #   orbax-checkpoint
+flax==0.9.0 \
+    --hash=sha256:12cd8f7162165ddd56877fb1cd9a4fcb47a31569e4c5343eeb59a36369fa2cfe \
+    --hash=sha256:8b7f361eed0f5324e81f9dc8d02ea53da5f993d7c2e37e7aa5b37d3f6331dd53
+    # via
+    #   dm-haiku
+    #   jax-md
+frozendict==2.4.4 \
+    --hash=sha256:07c3a5dee8bbb84cba770e273cdbf2c87c8e035903af8f781292d72583416801 \
+    --hash=sha256:12a342e439aef28ccec533f0253ea53d75fe9102bd6ea928ff530e76eac38906 \
+    --hash=sha256:1697793b5f62b416c0fc1d94638ec91ed3aa4ab277f6affa3a95216ecb3af170 \
+    --hash=sha256:199a4d32194f3afed6258de7e317054155bc9519252b568d9cfffde7e4d834e5 \
+    --hash=sha256:259528ba6b56fa051bc996f1c4d8b57e30d6dd3bc2f27441891b04babc4b5e73 \
+    --hash=sha256:2b70b431e3a72d410a2cdf1497b3aba2f553635e0c0f657ce311d841bf8273b6 \
+    --hash=sha256:2bd009cf4fc47972838a91e9b83654dc9a095dc4f2bb3a37c3f3124c8a364543 \
+    --hash=sha256:2d8536e068d6bf281f23fa835ac07747fb0f8851879dd189e9709f9567408b4d \
+    --hash=sha256:3148062675536724502c6344d7c485dd4667fdf7980ca9bd05e338ccc0c4471e \
+    --hash=sha256:3f7c031b26e4ee6a3f786ceb5e3abf1181c4ade92dce1f847da26ea2c96008c7 \
+    --hash=sha256:4297d694eb600efa429769125a6f910ec02b85606f22f178bafbee309e7d3ec7 \
+    --hash=sha256:4a59578d47b3949437519b5c39a016a6116b9e787bb19289e333faae81462e59 \
+    --hash=sha256:4ae8d05c8d0b6134bfb6bfb369d5fa0c4df21eabb5ca7f645af95fdc6689678e \
+    --hash=sha256:5d58d9a8d9e49662c6dafbea5e641f97decdb3d6ccd76e55e79818415362ba25 \
+    --hash=sha256:63aa49f1919af7d45fb8fd5dec4c0859bc09f46880bd6297c79bb2db2969b63d \
+    --hash=sha256:6874fec816b37b6eb5795b00e0574cba261bf59723e2de607a195d5edaff0786 \
+    --hash=sha256:6eb716e6a6d693c03b1d53280a1947716129f5ef9bcdd061db5c17dea44b80fe \
+    --hash=sha256:705efca8d74d3facbb6ace80ab3afdd28eb8a237bfb4063ed89996b024bc443d \
+    --hash=sha256:78c94991944dd33c5376f720228e5b252ee67faf3bac50ef381adc9e51e90d9d \
+    --hash=sha256:7f79c26dff10ce11dad3b3627c89bb2e87b9dd5958c2b24325f16a23019b8b94 \
+    --hash=sha256:7fee9420475bb6ff357000092aa9990c2f6182b2bab15764330f4ad7de2eae49 \
+    --hash=sha256:812ab17522ba13637826e65454115a914c2da538356e85f43ecea069813e4b33 \
+    --hash=sha256:85375ec6e979e6373bffb4f54576a68bf7497c350861d20686ccae38aab69c0a \
+    --hash=sha256:87ebcde21565a14fe039672c25550060d6f6d88cf1f339beac094c3b10004eb0 \
+    --hash=sha256:93a7b19afb429cbf99d56faf436b45ef2fa8fe9aca89c49eb1610c3bd85f1760 \
+    --hash=sha256:b3b967d5065872e27b06f785a80c0ed0a45d1f7c9b85223da05358e734d858ca \
+    --hash=sha256:c6bf9260018d653f3cab9bd147bd8592bf98a5c6e338be0491ced3c196c034a3 \
+    --hash=sha256:c8f92425686323a950337da4b75b4c17a3327b831df8c881df24038d560640d4 \
+    --hash=sha256:d13b4310db337f4d2103867c5a05090b22bc4d50ca842093779ef541ea9c9eea \
+    --hash=sha256:d9647563e76adb05b7cde2172403123380871360a114f546b4ae1704510801e5 \
+    --hash=sha256:dc2228874eacae390e63fd4f2bb513b3144066a977dc192163c9f6c7f6de6474 \
+    --hash=sha256:e1b941132d79ce72d562a13341d38fc217bc1ee24d8c35a20d754e79ff99e038 \
+    --hash=sha256:fefeb700bc7eb8b4c2dc48704e4221860d254c8989fb53488540bc44e44a1ac2
+    # via jax-md
+fsspec==2024.9.0 \
+    --hash=sha256:4b0afb90c2f21832df142f292649035d80b421f60a9e1c027802e5a0da2b04e8 \
+    --hash=sha256:a0947d552d8a6efa72cc2c730b12c41d043509156966cca4fb157b0f2a0c574b
+    # via etils
+humanize==4.11.0 \
+    --hash=sha256:b53caaec8532bcb2fff70c8826f904c35943f8cecaca29d272d9df38092736c0 \
+    --hash=sha256:e66f36020a2d5a974c504bd2555cf770621dbdbb6d82f94a6857c0b1ea2608be
+    # via orbax-checkpoint
 idna==3.8 \
     --hash=sha256:050b4e5baadcd44d760cedbd2b8e639f2ff89bbc7a5730fcc662954303377aac \
     --hash=sha256:d838c2c0ed6fced7693d5e8ab8e734d5f8fda53a039c0164afb0b82e771e3603
     # via requests
-jax==0.4.31 \
+importlib-resources==6.4.5 \
+    --hash=sha256:980862a1d16c9e147a59603677fa2aa5fd82b87f223b6cb870695bcfce830065 \
+    --hash=sha256:ac29d5f956f01d5e4bb63102a5a19957f1b9175e45649977264a1416783bb717
+    # via etils
+jax[cuda12-pip]==0.4.31 ; sys_platform == "linux" \
     --hash=sha256:5688703735133d0dc537e99a1d646198a49c9d472d4715fde4bd437c44151bd7 \
     --hash=sha256:fd2d470643a0073d822737f0788f71391656af7e62cc5b2e7995ee390ceac287
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   e3nn-jax
+    #   flax
+    #   jax-md
+    #   jraph
+    #   optax
+    #   orbax-checkpoint
 jax-cuda12-pjrt==0.4.31 \
     --hash=sha256:3e77d1cfebeca06517254eb568f082037e1a2aa3ed8f63c543492ad8ab5a1585 \
     --hash=sha256:8961abb381d893a3c2392ad76ab2067a81f8f2514f3b47d2da3ac24283293fe0
     # via jax-cuda12-plugin
-jax-cuda12-plugin==0.4.31 ; sys_platform == "linux" \
+jax-cuda12-plugin[with-cuda]==0.4.31 \
     --hash=sha256:146f26928ca719a0daa14bb9a9f5a5cbfa20211e76ea05b7bb534277a658a995 \
     --hash=sha256:5048acdf29755303b2887d948137fd82891b48cdbc086e67640ebb976457cfcc \
     --hash=sha256:5cfa46d4106f70c31944c9fafcfa7b04c6f9886a1b35df5477cde7d9c43eb8bd \
     --hash=sha256:7a179e5e80dd9890972d777d597f3c902c6876e42bcf2edcfe4f3ec5a610472e \
     --hash=sha256:a3727a332fbeac625ab6d5ae63d0ed9e62d1e0b5011f72130c9bc96e797395d5 \
     --hash=sha256:cdb6d0c4009438a6a6bd7997ab8a1194beda9ae7322b8d265eea9e551b6c2b4e
+    # via jax
+jax-md @ https://github.com/wsmoses/jax-md/archive/1188490610b95023f8a51166c3f6b92da31e78fe.tar.gz \
+    --hash=sha256:d84b17f0604986364818d5b320e0c6bfa8c898b30095f3f0dc43aeab141d2dde
     # via -r builddeps/test-requirements.txt
 jaxlib==0.4.31 \
     --hash=sha256:185fb615ab6bd95315fbcbd951d84e71f9835d603db8c03c91faee98ce95ff4d \
@@ -149,10 +246,32 @@ jaxlib==0.4.31 \
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   e3nn-jax
     #   jax
-libtpu-nightly==0.1.dev20240729 ; sys_platform == "linux" \
-    --hash=sha256:e3e4a4305e673d95cdd0076785e506234fc63b51124eeb3e301727f68b6b0126
-    # via -r builddeps/test-requirements.txt
+    #   jax-md
+    #   jraph
+    #   optax
+    #   orbax-checkpoint
+jmp==0.0.4 \
+    --hash=sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730 \
+    --hash=sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d
+    # via dm-haiku
+jraph==0.0.6.dev0 \
+    --hash=sha256:350fe37bf717f934f1f84fd3370a480b3178bfcb61dfa217c738971308c57625 \
+    --hash=sha256:c3ac3a0b224b344eb6d367e8bc312d95ea41bf825d01ea31b80dd8c22c0dd8b8
+    # via jax-md
+markdown-it-py==3.0.0 \
+    --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \
+    --hash=sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb
+    # via rich
+mdurl==0.1.2 \
+    --hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
+    --hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
+    # via markdown-it-py
+ml-collections==0.1.1 \
+    --hash=sha256:3fefcc72ec433aa1e5d32307a3e474bbb67f405be814ea52a2166bfc9dbe68cc
+    # via jax-md
 ml-dtypes==0.4.0 \
     --hash=sha256:03e7cda6ef164eed0abb31df69d2c00c3a5ab3e2610b6d4c42183a43329c72a5 \
     --hash=sha256:2bb83fd064db43e67e67d021e547698af4c8d5c6190f2e9b1c53c09f6ff5531d \
@@ -174,6 +293,83 @@ ml-dtypes==0.4.0 \
     # via
     #   jax
     #   jaxlib
+    #   tensorstore
+mpmath==1.3.0 \
+    --hash=sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f \
+    --hash=sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c
+    # via sympy
+msgpack==1.1.0 \
+    --hash=sha256:06f5fd2f6bb2a7914922d935d3b8bb4a7fff3a9a91cfce6d06c13bc42bec975b \
+    --hash=sha256:071603e2f0771c45ad9bc65719291c568d4edf120b44eb36324dcb02a13bfddf \
+    --hash=sha256:0907e1a7119b337971a689153665764adc34e89175f9a34793307d9def08e6ca \
+    --hash=sha256:0f92a83b84e7c0749e3f12821949d79485971f087604178026085f60ce109330 \
+    --hash=sha256:115a7af8ee9e8cddc10f87636767857e7e3717b7a2e97379dc2054712693e90f \
+    --hash=sha256:13599f8829cfbe0158f6456374e9eea9f44eee08076291771d8ae93eda56607f \
+    --hash=sha256:17fb65dd0bec285907f68b15734a993ad3fc94332b5bb21b0435846228de1f39 \
+    --hash=sha256:2137773500afa5494a61b1208619e3871f75f27b03bcfca7b3a7023284140247 \
+    --hash=sha256:3180065ec2abbe13a4ad37688b61b99d7f9e012a535b930e0e683ad6bc30155b \
+    --hash=sha256:398b713459fea610861c8a7b62a6fec1882759f308ae0795b5413ff6a160cf3c \
+    --hash=sha256:3d364a55082fb2a7416f6c63ae383fbd903adb5a6cf78c5b96cc6316dc1cedc7 \
+    --hash=sha256:3df7e6b05571b3814361e8464f9304c42d2196808e0119f55d0d3e62cd5ea044 \
+    --hash=sha256:41c991beebf175faf352fb940bf2af9ad1fb77fd25f38d9142053914947cdbf6 \
+    --hash=sha256:42f754515e0f683f9c79210a5d1cad631ec3d06cea5172214d2176a42e67e19b \
+    --hash=sha256:452aff037287acb1d70a804ffd022b21fa2bb7c46bee884dbc864cc9024128a0 \
+    --hash=sha256:4676e5be1b472909b2ee6356ff425ebedf5142427842aa06b4dfd5117d1ca8a2 \
+    --hash=sha256:46c34e99110762a76e3911fc923222472c9d681f1094096ac4102c18319e6468 \
+    --hash=sha256:471e27a5787a2e3f974ba023f9e265a8c7cfd373632247deb225617e3100a3c7 \
+    --hash=sha256:4a1964df7b81285d00a84da4e70cb1383f2e665e0f1f2a7027e683956d04b734 \
+    --hash=sha256:4b51405e36e075193bc051315dbf29168d6141ae2500ba8cd80a522964e31434 \
+    --hash=sha256:4d1b7ff2d6146e16e8bd665ac726a89c74163ef8cd39fa8c1087d4e52d3a2325 \
+    --hash=sha256:53258eeb7a80fc46f62fd59c876957a2d0e15e6449a9e71842b6d24419d88ca1 \
+    --hash=sha256:534480ee5690ab3cbed89d4c8971a5c631b69a8c0883ecfea96c19118510c846 \
+    --hash=sha256:58638690ebd0a06427c5fe1a227bb6b8b9fdc2bd07701bec13c2335c82131a88 \
+    --hash=sha256:58dfc47f8b102da61e8949708b3eafc3504509a5728f8b4ddef84bd9e16ad420 \
+    --hash=sha256:59caf6a4ed0d164055ccff8fe31eddc0ebc07cf7326a2aaa0dbf7a4001cd823e \
+    --hash=sha256:5dbad74103df937e1325cc4bfeaf57713be0b4f15e1c2da43ccdd836393e2ea2 \
+    --hash=sha256:5e1da8f11a3dd397f0a32c76165cf0c4eb95b31013a94f6ecc0b280c05c91b59 \
+    --hash=sha256:646afc8102935a388ffc3914b336d22d1c2d6209c773f3eb5dd4d6d3b6f8c1cb \
+    --hash=sha256:64fc9068d701233effd61b19efb1485587560b66fe57b3e50d29c5d78e7fef68 \
+    --hash=sha256:65553c9b6da8166e819a6aa90ad15288599b340f91d18f60b2061f402b9a4915 \
+    --hash=sha256:685ec345eefc757a7c8af44a3032734a739f8c45d1b0ac45efc5d8977aa4720f \
+    --hash=sha256:6ad622bf7756d5a497d5b6836e7fc3752e2dd6f4c648e24b1803f6048596f701 \
+    --hash=sha256:73322a6cc57fcee3c0c57c4463d828e9428275fb85a27aa2aa1a92fdc42afd7b \
+    --hash=sha256:74bed8f63f8f14d75eec75cf3d04ad581da6b914001b474a5d3cd3372c8cc27d \
+    --hash=sha256:79ec007767b9b56860e0372085f8504db5d06bd6a327a335449508bbee9648fa \
+    --hash=sha256:7a946a8992941fea80ed4beae6bff74ffd7ee129a90b4dd5cf9c476a30e9708d \
+    --hash=sha256:7ad442d527a7e358a469faf43fda45aaf4ac3249c8310a82f0ccff9164e5dccd \
+    --hash=sha256:7c9a35ce2c2573bada929e0b7b3576de647b0defbd25f5139dcdaba0ae35a4cc \
+    --hash=sha256:7e7b853bbc44fb03fbdba34feb4bd414322180135e2cb5164f20ce1c9795ee48 \
+    --hash=sha256:879a7b7b0ad82481c52d3c7eb99bf6f0645dbdec5134a4bddbd16f3506947feb \
+    --hash=sha256:8a706d1e74dd3dea05cb54580d9bd8b2880e9264856ce5068027eed09680aa74 \
+    --hash=sha256:8a84efb768fb968381e525eeeb3d92857e4985aacc39f3c47ffd00eb4509315b \
+    --hash=sha256:8cf9e8c3a2153934a23ac160cc4cba0ec035f6867c8013cc6077a79823370346 \
+    --hash=sha256:8da4bf6d54ceed70e8861f833f83ce0814a2b72102e890cbdfe4b34764cdd66e \
+    --hash=sha256:8e59bca908d9ca0de3dc8684f21ebf9a690fe47b6be93236eb40b99af28b6ea6 \
+    --hash=sha256:914571a2a5b4e7606997e169f64ce53a8b1e06f2cf2c3a7273aa106236d43dd5 \
+    --hash=sha256:a51abd48c6d8ac89e0cfd4fe177c61481aca2d5e7ba42044fd218cfd8ea9899f \
+    --hash=sha256:a52a1f3a5af7ba1c9ace055b659189f6c669cf3657095b50f9602af3a3ba0fe5 \
+    --hash=sha256:ad33e8400e4ec17ba782f7b9cf868977d867ed784a1f5f2ab46e7ba53b6e1e1b \
+    --hash=sha256:b4c01941fd2ff87c2a934ee6055bda4ed353a7846b8d4f341c428109e9fcde8c \
+    --hash=sha256:bce7d9e614a04d0883af0b3d4d501171fbfca038f12c77fa838d9f198147a23f \
+    --hash=sha256:c40ffa9a15d74e05ba1fe2681ea33b9caffd886675412612d93ab17b58ea2fec \
+    --hash=sha256:c5a91481a3cc573ac8c0d9aace09345d989dc4a0202b7fcb312c88c26d4e71a8 \
+    --hash=sha256:c921af52214dcbb75e6bdf6a661b23c3e6417f00c603dd2070bccb5c3ef499f5 \
+    --hash=sha256:d46cf9e3705ea9485687aa4001a76e44748b609d260af21c4ceea7f2212a501d \
+    --hash=sha256:d8ce0b22b890be5d252de90d0e0d119f363012027cf256185fc3d474c44b1b9e \
+    --hash=sha256:dd432ccc2c72b914e4cb77afce64aab761c1137cc698be3984eee260bcb2896e \
+    --hash=sha256:e0856a2b7e8dcb874be44fea031d22e5b3a19121be92a1e098f46068a11b0870 \
+    --hash=sha256:e1f3c3d21f7cf67bcf2da8e494d30a75e4cf60041d98b3f79875afb5b96f3a3f \
+    --hash=sha256:f1ba6136e650898082d9d5a5217d5906d1e138024f836ff48691784bbe1adf96 \
+    --hash=sha256:f3e9b4936df53b970513eac1758f3882c88658a220b58dcc1e39606dccaaf01c \
+    --hash=sha256:f80bc7d47f76089633763f952e67f8214cb7b3ee6bfa489b3cb6a84cfac114cd \
+    --hash=sha256:fd2906780f25c8ed5d7b323379f6138524ba793428db5d0e9d226d3fa6aa1788
+    # via
+    #   flax
+    #   orbax-checkpoint
+nest-asyncio==1.6.0 \
+    --hash=sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe \
+    --hash=sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c
+    # via orbax-checkpoint
 numpy==2.1.0 \
     --hash=sha256:08801848a40aea24ce16c2ecde3b756f9ad756586fb2d13210939eb69b023f5b \
     --hash=sha256:0937e54c09f7a9a68da6889362ddd2ff584c02d015ec92672c099b61555f8911 \
@@ -229,19 +425,175 @@ numpy==2.1.0 \
     --hash=sha256:ffbd6faeb190aaf2b5e9024bac9622d2ee549b7ec89ef3a9373fa35313d44e0e
     # via
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   dm-haiku
+    #   e3nn-jax
+    #   flax
     #   jax
+    #   jax-md
     #   jaxlib
+    #   jmp
+    #   jraph
     #   ml-dtypes
     #   opt-einsum
+    #   optax
+    #   orbax-checkpoint
     #   scipy
+    #   tensorstore
+nvidia-cublas-cu12==12.6.1.4 \
+    --hash=sha256:5dd125ece5469dbdceebe2e9536ad8fc4abd38aa394a7ace42fc8a930a1e81e3 \
+    --hash=sha256:5e5d384583d72ac364064ced3dd92a5caa59a8a57568595c9f82e83d255b2481 \
+    --hash=sha256:c25ab29026a265d46c1063b5fb3cb9440f5f2eb88041c6b7c6711bcb3361789f
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+nvidia-cuda-cupti-cu12==12.6.68 \
+    --hash=sha256:13408a021727de6473d138a0c5e8080b23437f761508e2b11d2530fed24f4ea0 \
+    --hash=sha256:5ad6a1fcfcb42c8628f7e547079575116d428d0cb3b4fab98362e08a9ea0b842 \
+    --hash=sha256:7487f59d73a090bf661fa8da84bad649f019a249dbac3a6cc58b039e15c28d91
+    # via jax-cuda12-plugin
+nvidia-cuda-nvcc-cu12==12.6.68 \
+    --hash=sha256:3999aa4a42ac8723c09a8aafd06bc4a6ec1a0b05c53bc96c8d6cf195e84f6935 \
+    --hash=sha256:9c0a18d76f0d1de99ba1d5fd70cffb32c0249e4abc42de9c0504e34d90ff421c \
+    --hash=sha256:d2faca18a3d5dd48865ad259262f7da43358d0940d53554026102d70c14ea2f9
+    # via jax-cuda12-plugin
+nvidia-cuda-runtime-cu12==12.6.68 \
+    --hash=sha256:3d421aa4ff608b2d8c650e0208a0fb28b4b6792a35b42bd2769d802149f85238 \
+    --hash=sha256:806b51a1dd266aac41ae09ca6142faee1686d119ced006cb9b76dfd331c75ab8 \
+    --hash=sha256:846987485889786d257f6d7bdcf7544a36452936514e20dd710527b896c0fe12
+    # via jax-cuda12-plugin
+nvidia-cudnn-cu12==9.3.0.75 \
+    --hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
+    --hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
+    --hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
+    # via jax-cuda12-plugin
+nvidia-cufft-cu12==11.2.6.59 \
+    --hash=sha256:251df5b20b11bb2af6d3964ac01b85a94094222d081c90f27e8df3bf533d3257 \
+    --hash=sha256:2ea19d2101d309228daeb1045397d8e28eb3ec1ec45f226bdc12ac6e9c1c59d4 \
+    --hash=sha256:998bbd77799dc427f9c48e5d57a316a7370d231fd96121fb018b370f67fc4909
+    # via jax-cuda12-plugin
+nvidia-cusolver-cu12==11.6.4.69 \
+    --hash=sha256:07d9a1fc00049cba615ec3475eca5320943df3175b05d358d2559286bb7f1fa6 \
+    --hash=sha256:1c799e473bbd369a34490322ebf6bbf8862831e199f5d6da6868d5f6f7332fff \
+    --hash=sha256:ec0419e653587d25f399736eaf1d26a6562d8bcaeb44b1e3daef87e13b669963
+    # via jax-cuda12-plugin
+nvidia-cusparse-cu12==12.5.3.3 \
+    --hash=sha256:76030755020d3a969b40273f43b8c496c4e122ee2a01fd724cf1398421bcadd8 \
+    --hash=sha256:bfa07cb86edfd6112dbead189c182a924fd9cb3e48ae117b1ac4cd3084078bc0 \
+    --hash=sha256:c9d0ff7870672b1e0c7ffc1e47e9b87b51e38ad32ae39e05f08fc68933983a80
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cusolver-cu12
+nvidia-nccl-cu12==2.22.3 \
+    --hash=sha256:f9f5e03c00269dee2cd1aa57019f9a024478a74ae6e9b32d5341c849fe6f6302
+    # via jax-cuda12-plugin
+nvidia-nvjitlink-cu12==12.6.68 \
+    --hash=sha256:125a6c2a44e96386dda634e13d944e60b07a0402d391a070e8fb4104b34ea1ab \
+    --hash=sha256:a55744c98d70317c5e23db14866a8cc2b733f7324509e941fc96276f9f37801d \
+    --hash=sha256:b3fd0779845f68b92063ab1393abab1ed0a23412fc520df79a8190d098b5cd6b
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cufft-cu12
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
 opt-einsum==3.3.0 \
     --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
     --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
     # via jax
+optax==0.2.3 \
+    --hash=sha256:083e603dcd731d7e74d99f71c12f77937dd53f79001b4c09c290e4f47dd2e94f \
+    --hash=sha256:ec7ab925440b0c5a512e1f24fba0fb3e7d760a7fd5d2496d7a691e9d37da01d9
+    # via
+    #   flax
+    #   jax-md
+orbax-checkpoint==0.6.4 \
+    --hash=sha256:366b4d528a7322e1b3d9ddcaed45c8515add0d2fc69c8975c30d98638543240f \
+    --hash=sha256:b4f2608ee4d1da67f7619fc35ff9c928ecdf4ccf7546eeb43ecf38c2608b6dea
+    # via flax
+protobuf==5.28.2 \
+    --hash=sha256:2c69461a7fcc8e24be697624c09a839976d82ae75062b11a0972e41fd2cd9132 \
+    --hash=sha256:35cfcb15f213449af7ff6198d6eb5f739c37d7e4f1c09b5d0641babf2cc0c68f \
+    --hash=sha256:52235802093bd8a2811abbe8bf0ab9c5f54cca0a751fdd3f6ac2a21438bffece \
+    --hash=sha256:59379674ff119717404f7454647913787034f03fe7049cbef1d74a97bb4593f0 \
+    --hash=sha256:5e8a95246d581eef20471b5d5ba010d55f66740942b95ba9b872d918c459452f \
+    --hash=sha256:87317e9bcda04a32f2ee82089a204d3a2f0d3c8aeed16568c7daf4756e4f1fe0 \
+    --hash=sha256:8ddc60bf374785fb7cb12510b267f59067fa10087325b8e1855b898a0d81d276 \
+    --hash=sha256:a8b9403fc70764b08d2f593ce44f1d2920c5077bf7d311fefec999f8c40f78b7 \
+    --hash=sha256:c0ea0123dac3399a2eeb1a1443d82b7afc9ff40241433296769f7da42d142ec3 \
+    --hash=sha256:ca53faf29896c526863366a52a8f4d88e69cd04ec9571ed6082fa117fac3ab36 \
+    --hash=sha256:eeea10f3dc0ac7e6b4933d32db20662902b4ab81bf28df12218aa389e9c2102d
+    # via orbax-checkpoint
+pygments==2.18.0 \
+    --hash=sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199 \
+    --hash=sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a
+    # via rich
+pyyaml==6.0.2 \
+    --hash=sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff \
+    --hash=sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48 \
+    --hash=sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086 \
+    --hash=sha256:0b69e4ce7a131fe56b7e4d770c67429700908fc0752af059838b1cfb41960e4e \
+    --hash=sha256:0ffe8360bab4910ef1b9e87fb812d8bc0a308b0d0eef8c8f44e0254ab3b07133 \
+    --hash=sha256:11d8f3dd2b9c1207dcaf2ee0bbbfd5991f571186ec9cc78427ba5bd32afae4b5 \
+    --hash=sha256:17e311b6c678207928d649faa7cb0d7b4c26a0ba73d41e99c4fff6b6c3276484 \
+    --hash=sha256:1e2120ef853f59c7419231f3bf4e7021f1b936f6ebd222406c3b60212205d2ee \
+    --hash=sha256:1f71ea527786de97d1a0cc0eacd1defc0985dcf6b3f17bb77dcfc8c34bec4dc5 \
+    --hash=sha256:23502f431948090f597378482b4812b0caae32c22213aecf3b55325e049a6c68 \
+    --hash=sha256:24471b829b3bf607e04e88d79542a9d48bb037c2267d7927a874e6c205ca7e9a \
+    --hash=sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf \
+    --hash=sha256:2e99c6826ffa974fe6e27cdb5ed0021786b03fc98e5ee3c5bfe1fd5015f42b99 \
+    --hash=sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8 \
+    --hash=sha256:3ad2a3decf9aaba3d29c8f537ac4b243e36bef957511b4766cb0057d32b0be85 \
+    --hash=sha256:3b1fdb9dc17f5a7677423d508ab4f243a726dea51fa5e70992e59a7411c89d19 \
+    --hash=sha256:41e4e3953a79407c794916fa277a82531dd93aad34e29c2a514c2c0c5fe971cc \
+    --hash=sha256:43fa96a3ca0d6b1812e01ced1044a003533c47f6ee8aca31724f78e93ccc089a \
+    --hash=sha256:50187695423ffe49e2deacb8cd10510bc361faac997de9efef88badc3bb9e2d1 \
+    --hash=sha256:5ac9328ec4831237bec75defaf839f7d4564be1e6b25ac710bd1a96321cc8317 \
+    --hash=sha256:5d225db5a45f21e78dd9358e58a98702a0302f2659a3c6cd320564b75b86f47c \
+    --hash=sha256:6395c297d42274772abc367baaa79683958044e5d3835486c16da75d2a694631 \
+    --hash=sha256:688ba32a1cffef67fd2e9398a2efebaea461578b0923624778664cc1c914db5d \
+    --hash=sha256:68ccc6023a3400877818152ad9a1033e3db8625d899c72eacb5a668902e4d652 \
+    --hash=sha256:70b189594dbe54f75ab3a1acec5f1e3faa7e8cf2f1e08d9b561cb41b845f69d5 \
+    --hash=sha256:797b4f722ffa07cc8d62053e4cff1486fa6dc094105d13fea7b1de7d8bf71c9e \
+    --hash=sha256:7c36280e6fb8385e520936c3cb3b8042851904eba0e58d277dca80a5cfed590b \
+    --hash=sha256:7e7401d0de89a9a855c839bc697c079a4af81cf878373abd7dc625847d25cbd8 \
+    --hash=sha256:80bab7bfc629882493af4aa31a4cfa43a4c57c83813253626916b8c7ada83476 \
+    --hash=sha256:82d09873e40955485746739bcb8b4586983670466c23382c19cffecbf1fd8706 \
+    --hash=sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563 \
+    --hash=sha256:8824b5a04a04a047e72eea5cec3bc266db09e35de6bdfe34c9436ac5ee27d237 \
+    --hash=sha256:8b9c7197f7cb2738065c481a0461e50ad02f18c78cd75775628afb4d7137fb3b \
+    --hash=sha256:9056c1ecd25795207ad294bcf39f2db3d845767be0ea6e6a34d856f006006083 \
+    --hash=sha256:936d68689298c36b53b29f23c6dbb74de12b4ac12ca6cfe0e047bedceea56180 \
+    --hash=sha256:9b22676e8097e9e22e36d6b7bda33190d0d400f345f23d4065d48f4ca7ae0425 \
+    --hash=sha256:a4d3091415f010369ae4ed1fc6b79def9416358877534caf6a0fdd2146c87a3e \
+    --hash=sha256:a8786accb172bd8afb8be14490a16625cbc387036876ab6ba70912730faf8e1f \
+    --hash=sha256:a9f8c2e67970f13b16084e04f134610fd1d374bf477b17ec1599185cf611d725 \
+    --hash=sha256:bc2fa7c6b47d6bc618dd7fb02ef6fdedb1090ec036abab80d4681424b84c1183 \
+    --hash=sha256:c70c95198c015b85feafc136515252a261a84561b7b1d51e3384e0655ddf25ab \
+    --hash=sha256:cc1c1159b3d456576af7a3e4d1ba7e6924cb39de8f67111c735f6fc832082774 \
+    --hash=sha256:ce826d6ef20b1bc864f0a68340c8b3287705cae2f8b4b1d932177dcc76721725 \
+    --hash=sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e \
+    --hash=sha256:d7fded462629cfa4b685c5416b949ebad6cec74af5e2d42905d41e257e0869f5 \
+    --hash=sha256:d84a1718ee396f54f3a086ea0a66d8e552b2ab2017ef8b420e92edbc841c352d \
+    --hash=sha256:d8e03406cac8513435335dbab54c0d385e4a49e4945d2909a581c83647ca0290 \
+    --hash=sha256:e10ce637b18caea04431ce14fabcf5c64a1c61ec9c56b071a4b7ca131ca52d44 \
+    --hash=sha256:ec031d5d2feb36d1d1a24380e4db6d43695f3748343d99434e6f5f9156aaa2ed \
+    --hash=sha256:ef6107725bd54b262d6dedcc2af448a266975032bc85ef0172c5f059da6325b4 \
+    --hash=sha256:efdca5630322a10774e8e98e1af481aad470dd62c3170801852d752aa7a783ba \
+    --hash=sha256:f753120cb8181e736c57ef7636e83f31b9c0d1722c516f7e86cf15b7aa57ff12 \
+    --hash=sha256:ff3824dc5261f50c9b0dfb3be22b4567a6f938ccce4587b38952d85fd9e9afe4
+    # via
+    #   flax
+    #   ml-collections
+    #   orbax-checkpoint
 requests==2.32.3 ; sys_platform == "linux" \
     --hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \
     --hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6
     # via -r builddeps/test-requirements.txt
+rich==13.9.2 \
+    --hash=sha256:51a2c62057461aaf7152b4d611168f93a9fc73068f8ded2790f29fe2b5366d0c \
+    --hash=sha256:8c82a3d3f8dcfe9e734771313e606b39d8247bb6b826e196f4914b333b743cf1
+    # via flax
 scipy==1.14.1 \
     --hash=sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e \
     --hash=sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79 \
@@ -279,7 +631,60 @@ scipy==1.14.1 \
     # via
     #   jax
     #   jaxlib
+six==1.16.0 \
+    --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
+    --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
+    # via ml-collections
+sympy==1.13.3 \
+    --hash=sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73 \
+    --hash=sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9
+    # via e3nn-jax
+tabulate==0.9.0 \
+    --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \
+    --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f
+    # via dm-haiku
+tensorstore==0.1.66 \
+    --hash=sha256:0261b1b7121186ac4898c9b4e29289c46754f9d7d73e906620e9a21386fe564e \
+    --hash=sha256:06f7dba630d75f65201b1c021fcd89e5f8538506b444b5dd712f35c11b812fc8 \
+    --hash=sha256:0d1f8c170d11e6e7a2132e3e0bb34565d24faff1fca42b5c7042772356e2fce0 \
+    --hash=sha256:0de07f34b181023f9179b1e2f7cb255dfad300a4d3f18639d5c4cc955ae88b6b \
+    --hash=sha256:16c4474b10c90b27d422ccffd31f78ed7c2f2976330ac80a4b53a73fd8f041d9 \
+    --hash=sha256:356e2a7fa256060c9da5d88cac6d9384d99887ab0df819c61f1dbe298ab56085 \
+    --hash=sha256:3bd37baa5983b9707c2b5e63301a2091aac3c21fef9891b5fce67597f9cfd61e \
+    --hash=sha256:486be8ec698d2eb10f297353351ff69e90551297185a158d6875e097fa08de68 \
+    --hash=sha256:5af8bf532ae4cc75709093483043ccd1d0f480d0f2b5e965ee5724f7d229658c \
+    --hash=sha256:5ca2a389b8b274f3a8bef907469bc07f700001ee8ae0bc63a5164135c5afa8ee \
+    --hash=sha256:6de7521a0f1f0fb4c4bb80ba945cdd926c4301d1ff2d821946f1675c0f532f5c \
+    --hash=sha256:94523e656c013721924cb5a2955b6c80f3bc4b469dd3d89cb759ab38f1c96cac \
+    --hash=sha256:9c7a672579b08a9c6af4875c56ced1f61f19df4cabe7b61b445a07307fe85d0a \
+    --hash=sha256:a150257280f580f73b3c037a8e60c56a51e5c81d4bc3a1a90cab0e00a292dddd \
+    --hash=sha256:a92dfc7b05e959e5823a87b8a20dc291382d5a647dff49eb560d5516e3490346 \
+    --hash=sha256:afcfa90c5dfc969e2286a30df97b91588486f77b232793abdd7c2cfd20df29d5 \
+    --hash=sha256:b77ee47da8a1b3d6fd03e23a8f853a2a666037f03e21546b4eb2b4cf43e13a96 \
+    --hash=sha256:bb3d9bc9156adaf2212953d941efd993b57a8da7053b93f512c1c313dfdddb89 \
+    --hash=sha256:bbc662c4697cd21443518ffb0c8f23aa72073c6a80f9828aea34a3b43d2230dc \
+    --hash=sha256:c4047f37ea0fbf665e25e918ad9d178fb2b3513203654ae10d8853d14e10a107 \
+    --hash=sha256:c974926fb2f3048fa3d36fc2fdf572403c7b70191193e83c5af78fdc46910b62
+    # via
+    #   flax
+    #   orbax-checkpoint
+toolz==1.0.0 \
+    --hash=sha256:292c8f1c4e7516bf9086f8850935c799a874039c8bcf959d47b600e4c44a6236 \
+    --hash=sha256:2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02
+    # via chex
+typing-extensions==4.12.2 \
+    --hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \
+    --hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8
+    # via
+    #   chex
+    #   etils
+    #   flax
+    #   orbax-checkpoint
 urllib3==2.2.2 \
     --hash=sha256:a448b2f64d686155468037e1ace9f2d2199776e17f0a46610480d311f73e3472 \
     --hash=sha256:dd505485549a7a552833da5e6063639d0d177c04f23bc3864e41e5dc5f612168
     # via requests
+zipp==3.20.2 \
+    --hash=sha256:a817ac80d6cf4b23bf7f2828b7cabf326f15a001bea8b1f9b49631780ba28350 \
+    --hash=sha256:bc9eb26f4506fda01b81bcde0ca78103b6e62f991b381fec825435c836edbc29
+    # via etils
diff --git a/builddeps/requirements_lock_3_12.txt b/builddeps/requirements_lock_3_12.txt
index a223e5db..41964ff9 100644
--- a/builddeps/requirements_lock_3_12.txt
+++ b/builddeps/requirements_lock_3_12.txt
@@ -4,14 +4,22 @@
 #
 #    bazel run //builddeps:requirements.update
 #
---find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
-
 absl-py==2.1.0 \
     --hash=sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308 \
     --hash=sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   dm-haiku
+    #   jax-md
+    #   ml-collections
+    #   optax
+    #   orbax-checkpoint
+attrs==24.2.0 \
+    --hash=sha256:5cfb1b9148b5b086569baec03f20d7b6bf3bcacc9a42bebf87ffaaca362f6346 \
+    --hash=sha256:81921eb96de3191c8258c199618104dd27ac608d9366f5e35d011eae1867ede2
+    # via e3nn-jax
 certifi==2024.7.4 \
     --hash=sha256:5a1e7645bc0ec61a09e26c36f6106dd4cf40c6db3a1fb6352b0244e7fb057c7b \
     --hash=sha256:c198e21b1289c2ab85ee4e67bb4b4ef3ead0892059901a8d5b622f24a1101e90
@@ -108,27 +116,116 @@ charset-normalizer==3.3.2 \
     --hash=sha256:fd1abc0d89e30cc4e02e4064dc67fcc51bd941eb395c502aac3ec19fab46b519 \
     --hash=sha256:ff8fa367d09b717b2a17a052544193ad76cd49979c805768879cb63d9ca50561
     # via requests
+chex==0.1.87 \
+    --hash=sha256:0096d89cc8d898bb521ef4bfbf5c24549022b0e5b301f529ab57238896fe6c5d \
+    --hash=sha256:ce536475661fd96d21be0c1728ecdbedd03f8ff950c662dfc338c92ea782cb16
+    # via optax
+contextlib2==21.6.0 \
+    --hash=sha256:3fbdb64466afd23abaf6c977627b75b6139a5a3e8ce38405c5b413aed7a0471f \
+    --hash=sha256:ab1e2bfe1d01d968e1b7e8d9023bc51ef3509bba217bb730cee3827e1ee82869
+    # via ml-collections
+dm-haiku==0.0.12 \
+    --hash=sha256:7448a43a6486bff95253f84e18eacc607d9c1256592573117a9d1d23e2780706 \
+    --hash=sha256:ba0b3acf71433156737fe342c486da11727e5e6c9e054245f4f9b8f0b53eb608
+    # via jax-md
+e3nn-jax==0.20.7 \
+    --hash=sha256:7764edb99d19465e574a23857e5f97420b4d720f06ba8bc18a01f9f2b8301d1c \
+    --hash=sha256:8a1b6f0c00feed565358fceec4c916a623412b6328945fce272f0c33e52b4bfc
+    # via jax-md
+einops==0.8.0 \
+    --hash=sha256:63486517fed345712a8385c100cb279108d9d47e6ae59099b07657e983deae85 \
+    --hash=sha256:9572fb63046264a862693b0a87088af3bdc8c068fde03de63453cbbde245465f
+    # via jax-md
+etils[epath,epy]==1.9.4 \
+    --hash=sha256:4387e7a4911a3b5cc4b92b99a9211386d176b43bae1dac8e2fe345fc2cb95e4b \
+    --hash=sha256:fad950414f0a1ca58c70c70915b0014f9953dd9bcf8aa951a0f75ff9becbeb24
+    # via
+    #   optax
+    #   orbax-checkpoint
+flax==0.9.0 \
+    --hash=sha256:12cd8f7162165ddd56877fb1cd9a4fcb47a31569e4c5343eeb59a36369fa2cfe \
+    --hash=sha256:8b7f361eed0f5324e81f9dc8d02ea53da5f993d7c2e37e7aa5b37d3f6331dd53
+    # via
+    #   dm-haiku
+    #   jax-md
+frozendict==2.4.4 \
+    --hash=sha256:07c3a5dee8bbb84cba770e273cdbf2c87c8e035903af8f781292d72583416801 \
+    --hash=sha256:12a342e439aef28ccec533f0253ea53d75fe9102bd6ea928ff530e76eac38906 \
+    --hash=sha256:1697793b5f62b416c0fc1d94638ec91ed3aa4ab277f6affa3a95216ecb3af170 \
+    --hash=sha256:199a4d32194f3afed6258de7e317054155bc9519252b568d9cfffde7e4d834e5 \
+    --hash=sha256:259528ba6b56fa051bc996f1c4d8b57e30d6dd3bc2f27441891b04babc4b5e73 \
+    --hash=sha256:2b70b431e3a72d410a2cdf1497b3aba2f553635e0c0f657ce311d841bf8273b6 \
+    --hash=sha256:2bd009cf4fc47972838a91e9b83654dc9a095dc4f2bb3a37c3f3124c8a364543 \
+    --hash=sha256:2d8536e068d6bf281f23fa835ac07747fb0f8851879dd189e9709f9567408b4d \
+    --hash=sha256:3148062675536724502c6344d7c485dd4667fdf7980ca9bd05e338ccc0c4471e \
+    --hash=sha256:3f7c031b26e4ee6a3f786ceb5e3abf1181c4ade92dce1f847da26ea2c96008c7 \
+    --hash=sha256:4297d694eb600efa429769125a6f910ec02b85606f22f178bafbee309e7d3ec7 \
+    --hash=sha256:4a59578d47b3949437519b5c39a016a6116b9e787bb19289e333faae81462e59 \
+    --hash=sha256:4ae8d05c8d0b6134bfb6bfb369d5fa0c4df21eabb5ca7f645af95fdc6689678e \
+    --hash=sha256:5d58d9a8d9e49662c6dafbea5e641f97decdb3d6ccd76e55e79818415362ba25 \
+    --hash=sha256:63aa49f1919af7d45fb8fd5dec4c0859bc09f46880bd6297c79bb2db2969b63d \
+    --hash=sha256:6874fec816b37b6eb5795b00e0574cba261bf59723e2de607a195d5edaff0786 \
+    --hash=sha256:6eb716e6a6d693c03b1d53280a1947716129f5ef9bcdd061db5c17dea44b80fe \
+    --hash=sha256:705efca8d74d3facbb6ace80ab3afdd28eb8a237bfb4063ed89996b024bc443d \
+    --hash=sha256:78c94991944dd33c5376f720228e5b252ee67faf3bac50ef381adc9e51e90d9d \
+    --hash=sha256:7f79c26dff10ce11dad3b3627c89bb2e87b9dd5958c2b24325f16a23019b8b94 \
+    --hash=sha256:7fee9420475bb6ff357000092aa9990c2f6182b2bab15764330f4ad7de2eae49 \
+    --hash=sha256:812ab17522ba13637826e65454115a914c2da538356e85f43ecea069813e4b33 \
+    --hash=sha256:85375ec6e979e6373bffb4f54576a68bf7497c350861d20686ccae38aab69c0a \
+    --hash=sha256:87ebcde21565a14fe039672c25550060d6f6d88cf1f339beac094c3b10004eb0 \
+    --hash=sha256:93a7b19afb429cbf99d56faf436b45ef2fa8fe9aca89c49eb1610c3bd85f1760 \
+    --hash=sha256:b3b967d5065872e27b06f785a80c0ed0a45d1f7c9b85223da05358e734d858ca \
+    --hash=sha256:c6bf9260018d653f3cab9bd147bd8592bf98a5c6e338be0491ced3c196c034a3 \
+    --hash=sha256:c8f92425686323a950337da4b75b4c17a3327b831df8c881df24038d560640d4 \
+    --hash=sha256:d13b4310db337f4d2103867c5a05090b22bc4d50ca842093779ef541ea9c9eea \
+    --hash=sha256:d9647563e76adb05b7cde2172403123380871360a114f546b4ae1704510801e5 \
+    --hash=sha256:dc2228874eacae390e63fd4f2bb513b3144066a977dc192163c9f6c7f6de6474 \
+    --hash=sha256:e1b941132d79ce72d562a13341d38fc217bc1ee24d8c35a20d754e79ff99e038 \
+    --hash=sha256:fefeb700bc7eb8b4c2dc48704e4221860d254c8989fb53488540bc44e44a1ac2
+    # via jax-md
+fsspec==2024.9.0 \
+    --hash=sha256:4b0afb90c2f21832df142f292649035d80b421f60a9e1c027802e5a0da2b04e8 \
+    --hash=sha256:a0947d552d8a6efa72cc2c730b12c41d043509156966cca4fb157b0f2a0c574b
+    # via etils
+humanize==4.11.0 \
+    --hash=sha256:b53caaec8532bcb2fff70c8826f904c35943f8cecaca29d272d9df38092736c0 \
+    --hash=sha256:e66f36020a2d5a974c504bd2555cf770621dbdbb6d82f94a6857c0b1ea2608be
+    # via orbax-checkpoint
 idna==3.8 \
     --hash=sha256:050b4e5baadcd44d760cedbd2b8e639f2ff89bbc7a5730fcc662954303377aac \
     --hash=sha256:d838c2c0ed6fced7693d5e8ab8e734d5f8fda53a039c0164afb0b82e771e3603
     # via requests
-jax==0.4.31 \
+importlib-resources==6.4.5 \
+    --hash=sha256:980862a1d16c9e147a59603677fa2aa5fd82b87f223b6cb870695bcfce830065 \
+    --hash=sha256:ac29d5f956f01d5e4bb63102a5a19957f1b9175e45649977264a1416783bb717
+    # via etils
+jax[cuda12-pip]==0.4.31 ; sys_platform == "linux" \
     --hash=sha256:5688703735133d0dc537e99a1d646198a49c9d472d4715fde4bd437c44151bd7 \
     --hash=sha256:fd2d470643a0073d822737f0788f71391656af7e62cc5b2e7995ee390ceac287
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   e3nn-jax
+    #   flax
+    #   jax-md
+    #   jraph
+    #   optax
+    #   orbax-checkpoint
 jax-cuda12-pjrt==0.4.31 \
     --hash=sha256:3e77d1cfebeca06517254eb568f082037e1a2aa3ed8f63c543492ad8ab5a1585 \
     --hash=sha256:8961abb381d893a3c2392ad76ab2067a81f8f2514f3b47d2da3ac24283293fe0
     # via jax-cuda12-plugin
-jax-cuda12-plugin==0.4.31 ; sys_platform == "linux" \
+jax-cuda12-plugin[with-cuda]==0.4.31 \
     --hash=sha256:146f26928ca719a0daa14bb9a9f5a5cbfa20211e76ea05b7bb534277a658a995 \
     --hash=sha256:5048acdf29755303b2887d948137fd82891b48cdbc086e67640ebb976457cfcc \
     --hash=sha256:5cfa46d4106f70c31944c9fafcfa7b04c6f9886a1b35df5477cde7d9c43eb8bd \
     --hash=sha256:7a179e5e80dd9890972d777d597f3c902c6876e42bcf2edcfe4f3ec5a610472e \
     --hash=sha256:a3727a332fbeac625ab6d5ae63d0ed9e62d1e0b5011f72130c9bc96e797395d5 \
     --hash=sha256:cdb6d0c4009438a6a6bd7997ab8a1194beda9ae7322b8d265eea9e551b6c2b4e
+    # via jax
+jax-md @ https://github.com/wsmoses/jax-md/archive/1188490610b95023f8a51166c3f6b92da31e78fe.tar.gz \
+    --hash=sha256:d84b17f0604986364818d5b320e0c6bfa8c898b30095f3f0dc43aeab141d2dde
     # via -r builddeps/test-requirements.txt
 jaxlib==0.4.31 \
     --hash=sha256:185fb615ab6bd95315fbcbd951d84e71f9835d603db8c03c91faee98ce95ff4d \
@@ -149,10 +246,32 @@ jaxlib==0.4.31 \
     # via
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   e3nn-jax
     #   jax
-libtpu-nightly==0.1.dev20240729 ; sys_platform == "linux" \
-    --hash=sha256:e3e4a4305e673d95cdd0076785e506234fc63b51124eeb3e301727f68b6b0126
-    # via -r builddeps/test-requirements.txt
+    #   jax-md
+    #   jraph
+    #   optax
+    #   orbax-checkpoint
+jmp==0.0.4 \
+    --hash=sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730 \
+    --hash=sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d
+    # via dm-haiku
+jraph==0.0.6.dev0 \
+    --hash=sha256:350fe37bf717f934f1f84fd3370a480b3178bfcb61dfa217c738971308c57625 \
+    --hash=sha256:c3ac3a0b224b344eb6d367e8bc312d95ea41bf825d01ea31b80dd8c22c0dd8b8
+    # via jax-md
+markdown-it-py==3.0.0 \
+    --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \
+    --hash=sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb
+    # via rich
+mdurl==0.1.2 \
+    --hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
+    --hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
+    # via markdown-it-py
+ml-collections==0.1.1 \
+    --hash=sha256:3fefcc72ec433aa1e5d32307a3e474bbb67f405be814ea52a2166bfc9dbe68cc
+    # via jax-md
 ml-dtypes==0.4.0 \
     --hash=sha256:03e7cda6ef164eed0abb31df69d2c00c3a5ab3e2610b6d4c42183a43329c72a5 \
     --hash=sha256:2bb83fd064db43e67e67d021e547698af4c8d5c6190f2e9b1c53c09f6ff5531d \
@@ -174,6 +293,83 @@ ml-dtypes==0.4.0 \
     # via
     #   jax
     #   jaxlib
+    #   tensorstore
+mpmath==1.3.0 \
+    --hash=sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f \
+    --hash=sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c
+    # via sympy
+msgpack==1.1.0 \
+    --hash=sha256:06f5fd2f6bb2a7914922d935d3b8bb4a7fff3a9a91cfce6d06c13bc42bec975b \
+    --hash=sha256:071603e2f0771c45ad9bc65719291c568d4edf120b44eb36324dcb02a13bfddf \
+    --hash=sha256:0907e1a7119b337971a689153665764adc34e89175f9a34793307d9def08e6ca \
+    --hash=sha256:0f92a83b84e7c0749e3f12821949d79485971f087604178026085f60ce109330 \
+    --hash=sha256:115a7af8ee9e8cddc10f87636767857e7e3717b7a2e97379dc2054712693e90f \
+    --hash=sha256:13599f8829cfbe0158f6456374e9eea9f44eee08076291771d8ae93eda56607f \
+    --hash=sha256:17fb65dd0bec285907f68b15734a993ad3fc94332b5bb21b0435846228de1f39 \
+    --hash=sha256:2137773500afa5494a61b1208619e3871f75f27b03bcfca7b3a7023284140247 \
+    --hash=sha256:3180065ec2abbe13a4ad37688b61b99d7f9e012a535b930e0e683ad6bc30155b \
+    --hash=sha256:398b713459fea610861c8a7b62a6fec1882759f308ae0795b5413ff6a160cf3c \
+    --hash=sha256:3d364a55082fb2a7416f6c63ae383fbd903adb5a6cf78c5b96cc6316dc1cedc7 \
+    --hash=sha256:3df7e6b05571b3814361e8464f9304c42d2196808e0119f55d0d3e62cd5ea044 \
+    --hash=sha256:41c991beebf175faf352fb940bf2af9ad1fb77fd25f38d9142053914947cdbf6 \
+    --hash=sha256:42f754515e0f683f9c79210a5d1cad631ec3d06cea5172214d2176a42e67e19b \
+    --hash=sha256:452aff037287acb1d70a804ffd022b21fa2bb7c46bee884dbc864cc9024128a0 \
+    --hash=sha256:4676e5be1b472909b2ee6356ff425ebedf5142427842aa06b4dfd5117d1ca8a2 \
+    --hash=sha256:46c34e99110762a76e3911fc923222472c9d681f1094096ac4102c18319e6468 \
+    --hash=sha256:471e27a5787a2e3f974ba023f9e265a8c7cfd373632247deb225617e3100a3c7 \
+    --hash=sha256:4a1964df7b81285d00a84da4e70cb1383f2e665e0f1f2a7027e683956d04b734 \
+    --hash=sha256:4b51405e36e075193bc051315dbf29168d6141ae2500ba8cd80a522964e31434 \
+    --hash=sha256:4d1b7ff2d6146e16e8bd665ac726a89c74163ef8cd39fa8c1087d4e52d3a2325 \
+    --hash=sha256:53258eeb7a80fc46f62fd59c876957a2d0e15e6449a9e71842b6d24419d88ca1 \
+    --hash=sha256:534480ee5690ab3cbed89d4c8971a5c631b69a8c0883ecfea96c19118510c846 \
+    --hash=sha256:58638690ebd0a06427c5fe1a227bb6b8b9fdc2bd07701bec13c2335c82131a88 \
+    --hash=sha256:58dfc47f8b102da61e8949708b3eafc3504509a5728f8b4ddef84bd9e16ad420 \
+    --hash=sha256:59caf6a4ed0d164055ccff8fe31eddc0ebc07cf7326a2aaa0dbf7a4001cd823e \
+    --hash=sha256:5dbad74103df937e1325cc4bfeaf57713be0b4f15e1c2da43ccdd836393e2ea2 \
+    --hash=sha256:5e1da8f11a3dd397f0a32c76165cf0c4eb95b31013a94f6ecc0b280c05c91b59 \
+    --hash=sha256:646afc8102935a388ffc3914b336d22d1c2d6209c773f3eb5dd4d6d3b6f8c1cb \
+    --hash=sha256:64fc9068d701233effd61b19efb1485587560b66fe57b3e50d29c5d78e7fef68 \
+    --hash=sha256:65553c9b6da8166e819a6aa90ad15288599b340f91d18f60b2061f402b9a4915 \
+    --hash=sha256:685ec345eefc757a7c8af44a3032734a739f8c45d1b0ac45efc5d8977aa4720f \
+    --hash=sha256:6ad622bf7756d5a497d5b6836e7fc3752e2dd6f4c648e24b1803f6048596f701 \
+    --hash=sha256:73322a6cc57fcee3c0c57c4463d828e9428275fb85a27aa2aa1a92fdc42afd7b \
+    --hash=sha256:74bed8f63f8f14d75eec75cf3d04ad581da6b914001b474a5d3cd3372c8cc27d \
+    --hash=sha256:79ec007767b9b56860e0372085f8504db5d06bd6a327a335449508bbee9648fa \
+    --hash=sha256:7a946a8992941fea80ed4beae6bff74ffd7ee129a90b4dd5cf9c476a30e9708d \
+    --hash=sha256:7ad442d527a7e358a469faf43fda45aaf4ac3249c8310a82f0ccff9164e5dccd \
+    --hash=sha256:7c9a35ce2c2573bada929e0b7b3576de647b0defbd25f5139dcdaba0ae35a4cc \
+    --hash=sha256:7e7b853bbc44fb03fbdba34feb4bd414322180135e2cb5164f20ce1c9795ee48 \
+    --hash=sha256:879a7b7b0ad82481c52d3c7eb99bf6f0645dbdec5134a4bddbd16f3506947feb \
+    --hash=sha256:8a706d1e74dd3dea05cb54580d9bd8b2880e9264856ce5068027eed09680aa74 \
+    --hash=sha256:8a84efb768fb968381e525eeeb3d92857e4985aacc39f3c47ffd00eb4509315b \
+    --hash=sha256:8cf9e8c3a2153934a23ac160cc4cba0ec035f6867c8013cc6077a79823370346 \
+    --hash=sha256:8da4bf6d54ceed70e8861f833f83ce0814a2b72102e890cbdfe4b34764cdd66e \
+    --hash=sha256:8e59bca908d9ca0de3dc8684f21ebf9a690fe47b6be93236eb40b99af28b6ea6 \
+    --hash=sha256:914571a2a5b4e7606997e169f64ce53a8b1e06f2cf2c3a7273aa106236d43dd5 \
+    --hash=sha256:a51abd48c6d8ac89e0cfd4fe177c61481aca2d5e7ba42044fd218cfd8ea9899f \
+    --hash=sha256:a52a1f3a5af7ba1c9ace055b659189f6c669cf3657095b50f9602af3a3ba0fe5 \
+    --hash=sha256:ad33e8400e4ec17ba782f7b9cf868977d867ed784a1f5f2ab46e7ba53b6e1e1b \
+    --hash=sha256:b4c01941fd2ff87c2a934ee6055bda4ed353a7846b8d4f341c428109e9fcde8c \
+    --hash=sha256:bce7d9e614a04d0883af0b3d4d501171fbfca038f12c77fa838d9f198147a23f \
+    --hash=sha256:c40ffa9a15d74e05ba1fe2681ea33b9caffd886675412612d93ab17b58ea2fec \
+    --hash=sha256:c5a91481a3cc573ac8c0d9aace09345d989dc4a0202b7fcb312c88c26d4e71a8 \
+    --hash=sha256:c921af52214dcbb75e6bdf6a661b23c3e6417f00c603dd2070bccb5c3ef499f5 \
+    --hash=sha256:d46cf9e3705ea9485687aa4001a76e44748b609d260af21c4ceea7f2212a501d \
+    --hash=sha256:d8ce0b22b890be5d252de90d0e0d119f363012027cf256185fc3d474c44b1b9e \
+    --hash=sha256:dd432ccc2c72b914e4cb77afce64aab761c1137cc698be3984eee260bcb2896e \
+    --hash=sha256:e0856a2b7e8dcb874be44fea031d22e5b3a19121be92a1e098f46068a11b0870 \
+    --hash=sha256:e1f3c3d21f7cf67bcf2da8e494d30a75e4cf60041d98b3f79875afb5b96f3a3f \
+    --hash=sha256:f1ba6136e650898082d9d5a5217d5906d1e138024f836ff48691784bbe1adf96 \
+    --hash=sha256:f3e9b4936df53b970513eac1758f3882c88658a220b58dcc1e39606dccaaf01c \
+    --hash=sha256:f80bc7d47f76089633763f952e67f8214cb7b3ee6bfa489b3cb6a84cfac114cd \
+    --hash=sha256:fd2906780f25c8ed5d7b323379f6138524ba793428db5d0e9d226d3fa6aa1788
+    # via
+    #   flax
+    #   orbax-checkpoint
+nest-asyncio==1.6.0 \
+    --hash=sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe \
+    --hash=sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c
+    # via orbax-checkpoint
 numpy==2.1.0 \
     --hash=sha256:08801848a40aea24ce16c2ecde3b756f9ad756586fb2d13210939eb69b023f5b \
     --hash=sha256:0937e54c09f7a9a68da6889362ddd2ff584c02d015ec92672c099b61555f8911 \
@@ -229,19 +425,176 @@ numpy==2.1.0 \
     --hash=sha256:ffbd6faeb190aaf2b5e9024bac9622d2ee549b7ec89ef3a9373fa35313d44e0e
     # via
     #   -r builddeps/test-requirements.txt
+    #   chex
+    #   dm-haiku
+    #   e3nn-jax
+    #   flax
     #   jax
+    #   jax-md
     #   jaxlib
+    #   jmp
+    #   jraph
     #   ml-dtypes
     #   opt-einsum
+    #   optax
+    #   orbax-checkpoint
     #   scipy
+    #   tensorstore
+nvidia-cublas-cu12==12.6.3.3 \
+    --hash=sha256:e1f70bee38b964eac1907293b336bceb24498a4243e61eaf91a52977c59aebc4 \
+    --hash=sha256:e531199ca4f1f764fb45bc1dde49a006f6765033f9c89c737e4553b9502ca1f5 \
+    --hash=sha256:f33fb68e101d99470c82d17f92a0dd9f74de2a21685c217f4716cdd63b1316eb
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+nvidia-cuda-cupti-cu12==12.6.80 \
+    --hash=sha256:358b4a1d35370353d52e12f0a7d1769fc01ff74a191689d3870b2123156184c4 \
+    --hash=sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73 \
+    --hash=sha256:bbe6ae76e83ce5251b56e8c8e61a964f757175682bbad058b170b136266ab00a
+    # via jax-cuda12-plugin
+nvidia-cuda-nvcc-cu12==12.6.77 \
+    --hash=sha256:26e118b4dae4ebd316ab6f5ee3f3c4cb929f8265f717a8e016ba57e254c0b4f1 \
+    --hash=sha256:9a2aad43b8f4e4be9c237ece04f64450df77871eeebffad5a5a872714a97ca77 \
+    --hash=sha256:b66bb5dd6b8ae62262586691977d3b4a425e91db61b4b9bc2f6b42bcd4154b96
+    # via jax-cuda12-plugin
+nvidia-cuda-runtime-cu12==12.6.77 \
+    --hash=sha256:86c58044c824bf3c173c49a2dbc7a6c8b53cb4e4dca50068be0bf64e9dab3f7f \
+    --hash=sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8 \
+    --hash=sha256:d461264ecb429c84c8879a7153499ddc7b19b5f8d84c204307491989a365588e
+    # via jax-cuda12-plugin
+nvidia-cudnn-cu12==9.4.0.58 \
+    --hash=sha256:2b891934e9e9bfcce61ce7f123ac3d9fd3a3d21ee9f46426aa7a96a6e9452410 \
+    --hash=sha256:cbbc57da0cbab1f7f3a9b7790f702b75c9adb00ee67499e84dba2b458065749b \
+    --hash=sha256:eabcbb94ebc0c8eb12e42133e4aaec6b2d73750beb144d360f7f60fec81489e1
+    # via jax-cuda12-plugin
+nvidia-cufft-cu12==11.3.0.4 \
+    --hash=sha256:6048ebddfb90d09d2707efb1fd78d4e3a77cb3ae4dc60e19aab6be0ece2ae464 \
+    --hash=sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca \
+    --hash=sha256:8510990de9f96c803a051822618d42bf6cb8f069ff3f48d93a8486efdacb48fb
+    # via jax-cuda12-plugin
+nvidia-cusolver-cu12==11.7.1.2 \
+    --hash=sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0 \
+    --hash=sha256:6813f9d8073f555444a8705f3ab0296d3e1cb37a16d694c5fc8b862a0d8706d7 \
+    --hash=sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6
+    # via jax-cuda12-plugin
+nvidia-cusparse-cu12==12.5.4.2 \
+    --hash=sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f \
+    --hash=sha256:4acb8c08855a26d737398cba8fb6f8f5045d93f82612b4cfd84645a2332ccf20 \
+    --hash=sha256:7aa32fa5470cf754f72d1116c7cbc300b4e638d3ae5304cfa4a638a5b87161b1
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cusolver-cu12
+nvidia-nccl-cu12==2.23.4 \
+    --hash=sha256:aa946c8327e22ced28e7cef508a334673abc42064ec85f02d005ba1785ea4cec \
+    --hash=sha256:b097258d9aab2fa9f686e33c6fe40ae57b27df60cedbd15d139701bb5509e0c1
+    # via jax-cuda12-plugin
+nvidia-nvjitlink-cu12==12.6.77 \
+    --hash=sha256:3bf10d85bb1801e9c894c6e197e44dd137d2a0a9e43f8450e9ad13f2df0dd52d \
+    --hash=sha256:410718cd44962bed862a31dd0318620f6f9a8b28a6291967bcfcb446a6516771 \
+    --hash=sha256:9ae346d16203ae4ea513be416495167a0101d33d2d14935aa9c1829a3fb45142
+    # via
+    #   jax-cuda12-plugin
+    #   nvidia-cufft-cu12
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
 opt-einsum==3.3.0 \
     --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
     --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
     # via jax
+optax==0.2.3 \
+    --hash=sha256:083e603dcd731d7e74d99f71c12f77937dd53f79001b4c09c290e4f47dd2e94f \
+    --hash=sha256:ec7ab925440b0c5a512e1f24fba0fb3e7d760a7fd5d2496d7a691e9d37da01d9
+    # via
+    #   flax
+    #   jax-md
+orbax-checkpoint==0.6.4 \
+    --hash=sha256:366b4d528a7322e1b3d9ddcaed45c8515add0d2fc69c8975c30d98638543240f \
+    --hash=sha256:b4f2608ee4d1da67f7619fc35ff9c928ecdf4ccf7546eeb43ecf38c2608b6dea
+    # via flax
+protobuf==5.28.2 \
+    --hash=sha256:2c69461a7fcc8e24be697624c09a839976d82ae75062b11a0972e41fd2cd9132 \
+    --hash=sha256:35cfcb15f213449af7ff6198d6eb5f739c37d7e4f1c09b5d0641babf2cc0c68f \
+    --hash=sha256:52235802093bd8a2811abbe8bf0ab9c5f54cca0a751fdd3f6ac2a21438bffece \
+    --hash=sha256:59379674ff119717404f7454647913787034f03fe7049cbef1d74a97bb4593f0 \
+    --hash=sha256:5e8a95246d581eef20471b5d5ba010d55f66740942b95ba9b872d918c459452f \
+    --hash=sha256:87317e9bcda04a32f2ee82089a204d3a2f0d3c8aeed16568c7daf4756e4f1fe0 \
+    --hash=sha256:8ddc60bf374785fb7cb12510b267f59067fa10087325b8e1855b898a0d81d276 \
+    --hash=sha256:a8b9403fc70764b08d2f593ce44f1d2920c5077bf7d311fefec999f8c40f78b7 \
+    --hash=sha256:c0ea0123dac3399a2eeb1a1443d82b7afc9ff40241433296769f7da42d142ec3 \
+    --hash=sha256:ca53faf29896c526863366a52a8f4d88e69cd04ec9571ed6082fa117fac3ab36 \
+    --hash=sha256:eeea10f3dc0ac7e6b4933d32db20662902b4ab81bf28df12218aa389e9c2102d
+    # via orbax-checkpoint
+pygments==2.18.0 \
+    --hash=sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199 \
+    --hash=sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a
+    # via rich
+pyyaml==6.0.2 \
+    --hash=sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff \
+    --hash=sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48 \
+    --hash=sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086 \
+    --hash=sha256:0b69e4ce7a131fe56b7e4d770c67429700908fc0752af059838b1cfb41960e4e \
+    --hash=sha256:0ffe8360bab4910ef1b9e87fb812d8bc0a308b0d0eef8c8f44e0254ab3b07133 \
+    --hash=sha256:11d8f3dd2b9c1207dcaf2ee0bbbfd5991f571186ec9cc78427ba5bd32afae4b5 \
+    --hash=sha256:17e311b6c678207928d649faa7cb0d7b4c26a0ba73d41e99c4fff6b6c3276484 \
+    --hash=sha256:1e2120ef853f59c7419231f3bf4e7021f1b936f6ebd222406c3b60212205d2ee \
+    --hash=sha256:1f71ea527786de97d1a0cc0eacd1defc0985dcf6b3f17bb77dcfc8c34bec4dc5 \
+    --hash=sha256:23502f431948090f597378482b4812b0caae32c22213aecf3b55325e049a6c68 \
+    --hash=sha256:24471b829b3bf607e04e88d79542a9d48bb037c2267d7927a874e6c205ca7e9a \
+    --hash=sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf \
+    --hash=sha256:2e99c6826ffa974fe6e27cdb5ed0021786b03fc98e5ee3c5bfe1fd5015f42b99 \
+    --hash=sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8 \
+    --hash=sha256:3ad2a3decf9aaba3d29c8f537ac4b243e36bef957511b4766cb0057d32b0be85 \
+    --hash=sha256:3b1fdb9dc17f5a7677423d508ab4f243a726dea51fa5e70992e59a7411c89d19 \
+    --hash=sha256:41e4e3953a79407c794916fa277a82531dd93aad34e29c2a514c2c0c5fe971cc \
+    --hash=sha256:43fa96a3ca0d6b1812e01ced1044a003533c47f6ee8aca31724f78e93ccc089a \
+    --hash=sha256:50187695423ffe49e2deacb8cd10510bc361faac997de9efef88badc3bb9e2d1 \
+    --hash=sha256:5ac9328ec4831237bec75defaf839f7d4564be1e6b25ac710bd1a96321cc8317 \
+    --hash=sha256:5d225db5a45f21e78dd9358e58a98702a0302f2659a3c6cd320564b75b86f47c \
+    --hash=sha256:6395c297d42274772abc367baaa79683958044e5d3835486c16da75d2a694631 \
+    --hash=sha256:688ba32a1cffef67fd2e9398a2efebaea461578b0923624778664cc1c914db5d \
+    --hash=sha256:68ccc6023a3400877818152ad9a1033e3db8625d899c72eacb5a668902e4d652 \
+    --hash=sha256:70b189594dbe54f75ab3a1acec5f1e3faa7e8cf2f1e08d9b561cb41b845f69d5 \
+    --hash=sha256:797b4f722ffa07cc8d62053e4cff1486fa6dc094105d13fea7b1de7d8bf71c9e \
+    --hash=sha256:7c36280e6fb8385e520936c3cb3b8042851904eba0e58d277dca80a5cfed590b \
+    --hash=sha256:7e7401d0de89a9a855c839bc697c079a4af81cf878373abd7dc625847d25cbd8 \
+    --hash=sha256:80bab7bfc629882493af4aa31a4cfa43a4c57c83813253626916b8c7ada83476 \
+    --hash=sha256:82d09873e40955485746739bcb8b4586983670466c23382c19cffecbf1fd8706 \
+    --hash=sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563 \
+    --hash=sha256:8824b5a04a04a047e72eea5cec3bc266db09e35de6bdfe34c9436ac5ee27d237 \
+    --hash=sha256:8b9c7197f7cb2738065c481a0461e50ad02f18c78cd75775628afb4d7137fb3b \
+    --hash=sha256:9056c1ecd25795207ad294bcf39f2db3d845767be0ea6e6a34d856f006006083 \
+    --hash=sha256:936d68689298c36b53b29f23c6dbb74de12b4ac12ca6cfe0e047bedceea56180 \
+    --hash=sha256:9b22676e8097e9e22e36d6b7bda33190d0d400f345f23d4065d48f4ca7ae0425 \
+    --hash=sha256:a4d3091415f010369ae4ed1fc6b79def9416358877534caf6a0fdd2146c87a3e \
+    --hash=sha256:a8786accb172bd8afb8be14490a16625cbc387036876ab6ba70912730faf8e1f \
+    --hash=sha256:a9f8c2e67970f13b16084e04f134610fd1d374bf477b17ec1599185cf611d725 \
+    --hash=sha256:bc2fa7c6b47d6bc618dd7fb02ef6fdedb1090ec036abab80d4681424b84c1183 \
+    --hash=sha256:c70c95198c015b85feafc136515252a261a84561b7b1d51e3384e0655ddf25ab \
+    --hash=sha256:cc1c1159b3d456576af7a3e4d1ba7e6924cb39de8f67111c735f6fc832082774 \
+    --hash=sha256:ce826d6ef20b1bc864f0a68340c8b3287705cae2f8b4b1d932177dcc76721725 \
+    --hash=sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e \
+    --hash=sha256:d7fded462629cfa4b685c5416b949ebad6cec74af5e2d42905d41e257e0869f5 \
+    --hash=sha256:d84a1718ee396f54f3a086ea0a66d8e552b2ab2017ef8b420e92edbc841c352d \
+    --hash=sha256:d8e03406cac8513435335dbab54c0d385e4a49e4945d2909a581c83647ca0290 \
+    --hash=sha256:e10ce637b18caea04431ce14fabcf5c64a1c61ec9c56b071a4b7ca131ca52d44 \
+    --hash=sha256:ec031d5d2feb36d1d1a24380e4db6d43695f3748343d99434e6f5f9156aaa2ed \
+    --hash=sha256:ef6107725bd54b262d6dedcc2af448a266975032bc85ef0172c5f059da6325b4 \
+    --hash=sha256:efdca5630322a10774e8e98e1af481aad470dd62c3170801852d752aa7a783ba \
+    --hash=sha256:f753120cb8181e736c57ef7636e83f31b9c0d1722c516f7e86cf15b7aa57ff12 \
+    --hash=sha256:ff3824dc5261f50c9b0dfb3be22b4567a6f938ccce4587b38952d85fd9e9afe4
+    # via
+    #   flax
+    #   ml-collections
+    #   orbax-checkpoint
 requests==2.32.3 ; sys_platform == "linux" \
     --hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \
     --hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6
     # via -r builddeps/test-requirements.txt
+rich==13.9.2 \
+    --hash=sha256:51a2c62057461aaf7152b4d611168f93a9fc73068f8ded2790f29fe2b5366d0c \
+    --hash=sha256:8c82a3d3f8dcfe9e734771313e606b39d8247bb6b826e196f4914b333b743cf1
+    # via flax
 scipy==1.14.1 \
     --hash=sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e \
     --hash=sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79 \
@@ -279,7 +632,66 @@ scipy==1.14.1 \
     # via
     #   jax
     #   jaxlib
+six==1.16.0 \
+    --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
+    --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
+    # via ml-collections
+sympy==1.13.3 \
+    --hash=sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73 \
+    --hash=sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9
+    # via e3nn-jax
+tabulate==0.9.0 \
+    --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \
+    --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f
+    # via dm-haiku
+tensorstore==0.1.66 \
+    --hash=sha256:0261b1b7121186ac4898c9b4e29289c46754f9d7d73e906620e9a21386fe564e \
+    --hash=sha256:06f7dba630d75f65201b1c021fcd89e5f8538506b444b5dd712f35c11b812fc8 \
+    --hash=sha256:0d1f8c170d11e6e7a2132e3e0bb34565d24faff1fca42b5c7042772356e2fce0 \
+    --hash=sha256:0de07f34b181023f9179b1e2f7cb255dfad300a4d3f18639d5c4cc955ae88b6b \
+    --hash=sha256:16c4474b10c90b27d422ccffd31f78ed7c2f2976330ac80a4b53a73fd8f041d9 \
+    --hash=sha256:356e2a7fa256060c9da5d88cac6d9384d99887ab0df819c61f1dbe298ab56085 \
+    --hash=sha256:3bd37baa5983b9707c2b5e63301a2091aac3c21fef9891b5fce67597f9cfd61e \
+    --hash=sha256:486be8ec698d2eb10f297353351ff69e90551297185a158d6875e097fa08de68 \
+    --hash=sha256:5af8bf532ae4cc75709093483043ccd1d0f480d0f2b5e965ee5724f7d229658c \
+    --hash=sha256:5ca2a389b8b274f3a8bef907469bc07f700001ee8ae0bc63a5164135c5afa8ee \
+    --hash=sha256:6de7521a0f1f0fb4c4bb80ba945cdd926c4301d1ff2d821946f1675c0f532f5c \
+    --hash=sha256:94523e656c013721924cb5a2955b6c80f3bc4b469dd3d89cb759ab38f1c96cac \
+    --hash=sha256:9c7a672579b08a9c6af4875c56ced1f61f19df4cabe7b61b445a07307fe85d0a \
+    --hash=sha256:a150257280f580f73b3c037a8e60c56a51e5c81d4bc3a1a90cab0e00a292dddd \
+    --hash=sha256:a92dfc7b05e959e5823a87b8a20dc291382d5a647dff49eb560d5516e3490346 \
+    --hash=sha256:afcfa90c5dfc969e2286a30df97b91588486f77b232793abdd7c2cfd20df29d5 \
+    --hash=sha256:b77ee47da8a1b3d6fd03e23a8f853a2a666037f03e21546b4eb2b4cf43e13a96 \
+    --hash=sha256:bb3d9bc9156adaf2212953d941efd993b57a8da7053b93f512c1c313dfdddb89 \
+    --hash=sha256:bbc662c4697cd21443518ffb0c8f23aa72073c6a80f9828aea34a3b43d2230dc \
+    --hash=sha256:c4047f37ea0fbf665e25e918ad9d178fb2b3513203654ae10d8853d14e10a107 \
+    --hash=sha256:c974926fb2f3048fa3d36fc2fdf572403c7b70191193e83c5af78fdc46910b62
+    # via
+    #   flax
+    #   orbax-checkpoint
+toolz==1.0.0 \
+    --hash=sha256:292c8f1c4e7516bf9086f8850935c799a874039c8bcf959d47b600e4c44a6236 \
+    --hash=sha256:2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02
+    # via chex
+typing-extensions==4.12.2 \
+    --hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \
+    --hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8
+    # via
+    #   chex
+    #   etils
+    #   flax
+    #   orbax-checkpoint
 urllib3==2.2.2 \
     --hash=sha256:a448b2f64d686155468037e1ace9f2d2199776e17f0a46610480d311f73e3472 \
     --hash=sha256:dd505485549a7a552833da5e6063639d0d177c04f23bc3864e41e5dc5f612168
     # via requests
+zipp==3.20.2 \
+    --hash=sha256:a817ac80d6cf4b23bf7f2828b7cabf326f15a001bea8b1f9b49631780ba28350 \
+    --hash=sha256:bc9eb26f4506fda01b81bcde0ca78103b6e62f991b381fec825435c836edbc29
+    # via etils
+
+# The following packages are considered to be unsafe in a requirements file:
+setuptools==75.1.0 \
+    --hash=sha256:35ab7fd3bcd95e6b7fd704e4a1539513edad446c097797f2985e0e4b960772f2 \
+    --hash=sha256:d59a21b17a275fb872a9c3dae73963160ae079f1049ed956880cd7c09b120538
+    # via chex
diff --git a/src/enzyme_ad/jax/BUILD b/src/enzyme_ad/jax/BUILD
index c51a2591..f5978b72 100644
--- a/src/enzyme_ad/jax/BUILD
+++ b/src/enzyme_ad/jax/BUILD
@@ -5,7 +5,7 @@ load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_library")
 load("@tensat_crate_index//:defs.bzl", "all_crate_deps")
 load("deps/rust_cxx_bridge.bzl", "rust_cxx_bridge")
 
-exports_files(["enzymexlamlir-opt.cpp", "deps/tensat/converted.txt", "deps/tensat/extractor/extract.py"])
+exports_files(["enzymexlamlir-opt.cpp", "deps/tensat/converted.txt", "deps/tensat/converted_multi.txt", "deps/tensat/extractor/extract.py"])
 licenses(["notice"])
 
 package(
@@ -47,6 +47,7 @@ py_library(
     srcs = [
         "__init__.py",
         "primitives.py",
+        "utils.py",
     ],
     visibility = ["//visibility:public"],
 )
@@ -282,7 +283,6 @@ cc_library(
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:TransformDialect",
         "@llvm-project//mlir:TransformDialectInterfaces",
-        "@llvm-project//mlir:Transforms",
         "@llvm-project//mlir:TransformUtils",
 
         "@stablehlo//:reference_ops",
@@ -311,9 +311,9 @@ cc_library(
         "@xla//xla/python/ifrt/hlo:hlo_program",
         "@xla//xla/ffi:call_frame",
         "@com_google_protobuf//:protobuf",
-        "@tsl//tsl/profiler/backends/cpu:annotation_stack_impl",
-        "@tsl//tsl/profiler/backends/cpu:traceme_recorder_impl",
-        "@tsl//tsl/profiler/utils:time_utils_impl",
+        "@xla//xla/tsl/profiler/backends/cpu:annotation_stack_impl",
+        "@xla//xla/tsl/profiler/backends/cpu:traceme_recorder_impl",
+        "@xla//xla/tsl/profiler/utils:time_utils_impl",
         "@tsl//tsl/platform:env_impl",
         "@xla//xla/stream_executor:stream_executor_impl",
         "@xla//xla/mlir/utils:type_util",
@@ -321,18 +321,27 @@ cc_library(
         "@stablehlo//:chlo_capi_objects",
         "@com_google_absl//absl/hash:hash",
         "@com_google_absl//absl/log:initialize",
-        "@llvm-project//mlir:CAPIIRObjects",
+        "@llvm-project//mlir:CAPIIRObjects",  
+
+        "@xla//xla/translate/stablehlo_to_hlo:translate",
+        "@xla//xla/translate/mhlo_to_hlo:mlir_hlo_to_hlo",
+
+        "@xla//xla/service/gpu/model:gpu_performance_model",
+        "@xla//xla/service/gpu/model:hlo_op_profile_proto_cc_impl",
+        "@xla//xla/service/gpu/model:gpu_hlo_cost_analysis",
+        "@xla//xla/service/gpu:gpu_latency_hiding_scheduler",
+        "@xla//xla/stream_executor:device_description",
     ] + select({
     "@bazel_tools//src/conditions:darwin": [],
     "//conditions:default": [
         "@xla//xla/stream_executor/cuda:all_runtime", 
-        "@xla//xla/stream_executor/rocm:all_runtime",
+        # "@xla//xla/stream_executor/rocm:all_runtime",
         "@xla//xla/service/gpu/model:hlo_op_profiles",
-        "@xla//xla/service/gpu/model:hlo_op_profile_proto_cc_impl",
         "@xla//xla/service/gpu:nvptx_compiler",
         # "@xla//xla/service/gpu:amdgpu_compiler",
         "@xla//xla/service/gpu:gpu_transfer_manager",
-        "@xla//xla/stream_executor:kernel", 
+        "@xla//xla/stream_executor:kernel",
+        "@xla//xla/stream_executor/gpu:redzone_allocator",
     ]}),
     alwayslink = 1,
     linkstatic = 1
@@ -355,14 +364,14 @@ pybind_library(
         "@tsl//tsl/platform:env",
         "@tsl//tsl/platform:env_impl",
         "@tsl//tsl/platform:tensor_float_32_utils",
-        "@tsl//tsl/profiler/backends/cpu:annotation_stack_impl",
-        "@tsl//tsl/profiler/backends/cpu:traceme_recorder_impl",
-        "@tsl//tsl/profiler/utils:time_utils_impl",
-        "@tsl//tsl/protobuf:dnn_proto_cc_impl",
-        "@tsl//tsl/protobuf:histogram_proto_cc",
-        "@tsl//tsl/protobuf:histogram_proto_cc_impl",
-        "@tsl//tsl/protobuf:protos_all_cc_impl",
- 
+        "@xla//xla/tsl/profiler/backends/cpu:annotation_stack_impl",
+        "@xla//xla/tsl/profiler/backends/cpu:traceme_recorder_impl",
+        "@xla//xla/tsl/profiler/utils:time_utils_impl",
+        "@xla//xla/tsl/protobuf:dnn_proto_cc_impl",
+        "@xla//xla/tsl/protobuf:histogram_proto_cc",
+        "@xla//xla/tsl/protobuf:histogram_proto_cc_impl",
+        "@xla//xla/tsl/protobuf:protos_all_cc_impl",
+
         # This is similar to xla_binary rule and is needed to make XLA client compile.
         "@xla//xla:autotune_results_proto_cc",
         "@xla//xla:autotune_results_proto_cc_impl",
@@ -387,6 +396,7 @@ pybind_library(
         "@xla//xla/service/memory_space_assignment:memory_space_assignment_proto_cc_impl",
         "@xla//xla/stream_executor:device_description_proto_cc",
         "@xla//xla/stream_executor:device_description_proto_cc_impl",
+        "@xla//xla/stream_executor/gpu:redzone_allocator",
         "@xla//xla/translate/mhlo_to_hlo:type_to_shape",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla:xla_data_proto_cc_impl",
diff --git a/src/enzyme_ad/jax/Implementations/HLODerivatives.td b/src/enzyme_ad/jax/Implementations/HLODerivatives.td
index 08862b81..7f05d1b7 100644
--- a/src/enzyme_ad/jax/Implementations/HLODerivatives.td
+++ b/src/enzyme_ad/jax/Implementations/HLODerivatives.td
@@ -456,6 +456,7 @@ def GradDataConvBatchGroupCount : GlobalExpr</*needsprimal*/0, /*needsshadow*/0,
 }]>;
 
 /// Helpers - Max/Min
+def GE : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "ComparisonDirection::GE">;
 def LT : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "ComparisonDirection::LT">;
 def GT : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "ComparisonDirection::GT">;
 
@@ -531,6 +532,8 @@ def : HLODerivative<"Atan2Op", (Op $x, $y),
 
 def : HLOReadOnlyIdentityOp<"BroadcastInDimOp">;
 
+def : HLOReadOnlyIdentityOp<"GatherOp", [0]>;
+
 def : HLOInactiveOp<"CeilOp">;
 
 def : HLODerivative<"CbrtOp", (Op $x), [
@@ -546,6 +549,9 @@ def : HLODerivative<"ComplexOp", (Op $re, $im), [
 
 def : HLOMemoryIdentityOp<"ConcatenateOp", [], [-1]>;
 
+def : HLOMemoryIdentityOp<"DynamicSliceOp", [], [0]>;
+def : HLOMemoryIdentityOp<"DynamicUpdateSliceOp", [], [0, 1]>;
+
 def : HLODerivative<"ConvolutionOp", (Op $lhs, $rhs),
                     [
                         (Convolution (TypeOf $lhs), (DiffeRet), $lhs, (GradFilterConvWindowStrides), (GradFilterConvPadding), (GradFilterConvLhsDialation), (GradFilterConvRhsDialation), (GradFilterConvWindowReversal), (GradFilterConvDimensionNumbers), (GradFilterConvFeatureGroupCount), (GradFilterConvBatchGroupCount), (ResultDotPrec)),
@@ -645,9 +651,18 @@ def : HLODerivative<"PowOp", (Op $x, $y),
                   ]
                   >;
 
-def : HLOReadOnlyIdentityOp<"ReshapeOp", [0], (Op $x), [(Reshape (TypeOf $x), (DiffeRet))]>;
+def : HLODerivative<"AbsOp", (Op $x), [
+  (Select (Compare $x, (HLOConstantFP<"0"> $x), (GE)), (DiffeRet), (Neg (DiffeRet)))
+]>;
 
-def : HLORegionTerminatorOp<"ReturnOp">;
+def : HLODerivative<"RemOp", (Op $x, $y), 
+                  [
+	            (DiffeRet),
+                    (CheckedMul (DiffeRet), (Neg (Mul (Sign (Div $x, $y):$div),(Floor (Abs $div)))))
+                  ]
+                  >;
+
+def : HLOReadOnlyIdentityOp<"ReshapeOp", [0], (Op $x), [(Reshape (TypeOf $x), (DiffeRet))]>;
 
 def : HLODerivative<"ReverseOp", (Op $x), [(Reverse (DiffeRet), (Dimensions))], (Reverse (Shadow $x), (Dimensions))>;
 
@@ -666,6 +681,13 @@ def : HLODerivative<"RsqrtOp", (Op $x),
                     ]
                   >;
 
+def : HLODerivative<"ConvertOp", (Op $x),
+                    [
+                        (Convert (TypeOf $x), (DiffeRet))
+                    ],
+                    (Convert (ResultTypes), (Shadow $x))
+                  >;
+
 def : HLODerivative<"SelectOp", (Op $cond, $lhs, $rhs),
                     [
                         (AssertingInactiveArg),
diff --git a/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp b/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
index 94ce0706..d5228bf5 100644
--- a/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
+++ b/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
@@ -176,8 +176,6 @@ class AutoDiffReduceFwd
     }
     Operation *shadow = builder.clone(*orig, map);
 
-    Value shadowRes = shadow->getResult(0);
-
     auto invAdd = gutils->invertedPointers.lookup(innerOp.getResult(0));
     gutils->invertedPointers.erase(innerOp.getResult(0));
     gutils->erase(invAdd.getDefiningOp());
@@ -189,12 +187,61 @@ class AutoDiffReduceFwd
     }
     cast<OpTy>(primal).getBody().front().eraseArguments(baToErase);
 
-    gutils->setDiffe(orig->getResult(0), shadowRes, builder);
+    size_t shadowCnt = 0;
+    for (auto origRes : orig->getResults()) {
+      if (!gutils->isConstantValue(origRes)) {
+        gutils->setDiffe(origRes, shadow->getResult(shadowCnt), builder);
+        shadowCnt++;
+      }
+    }
     gutils->eraseIfUnused(orig);
     return success();
   }
 };
 
+class AutoDiffWhileFwd
+    : public AutoDiffOpInterface::ExternalModel<AutoDiffWhileFwd, WhileOp> {
+public:
+  LogicalResult createForwardModeTangent(Operation *orig, OpBuilder &builder,
+                                         MGradientUtils *gutils) const {
+    auto op = cast<WhileOp>(orig);
+
+    llvm::SmallDenseSet<unsigned> operandPositionsToShadow;
+    llvm::SmallDenseSet<unsigned> resultPositionsToShadow;
+
+    for (auto &&[res, arg] :
+         llvm::zip(op->getResults(), op.getBody().front().getArguments()))
+      if (!gutils->isConstantValue(res) || !gutils->isConstantValue(arg)) {
+        operandPositionsToShadow.insert(res.getResultNumber());
+        resultPositionsToShadow.insert(res.getResultNumber());
+      }
+
+    auto res = mlir::enzyme::detail::controlFlowForwardHandler(
+        op, builder, gutils, operandPositionsToShadow, resultPositionsToShadow);
+
+    // Rewrite block arguments to match the shadowing
+    for (auto reg : {&op.getCond(), &op.getBody()}) {
+      auto nb = gutils->getNewFromOriginal(&reg->front());
+      size_t curidx = 0;
+      for (auto arg : reg->front().getArguments()) {
+        curidx++;
+        auto idx = arg.getArgNumber();
+        if (resultPositionsToShadow.count(idx)) {
+          if (gutils->isConstantValue(arg)) {
+            nb->insertArgument(
+                curidx,
+                cast<AutoDiffTypeInterface>(arg.getType()).getShadowType(),
+                op.getLoc());
+          }
+          curidx++;
+        }
+      }
+    }
+
+    return res;
+  }
+};
+
 class AutoDiffBroadcastInDimRev
     : public ReverseAutoDiffOpInterface::ExternalModel<
           AutoDiffBroadcastInDimRev, BroadcastInDimOp> {
@@ -499,6 +546,442 @@ struct SHLOConstantOpBatchInterface
   }
 };
 
+struct SHLOTransposeOpBatchInterface
+    : public BatchOpInterface::ExternalModel<SHLOTransposeOpBatchInterface,
+                                             TransposeOp> {
+
+  mlir::Operation *createBatch(Operation *src, IRMapping &mapper,
+                               Operation::CloneOptions options,
+                               std::map<Operation *, Operation *> &opMap,
+                               ArrayRef<int64_t> batchSizes) const {
+    SmallVector<Type> resultTypes(src->getResultTypes().begin(),
+                                  src->getResultTypes().end());
+    for (auto &Ty : resultTypes) {
+      auto T = cast<TensorType>(Ty);
+      SmallVector<int64_t> shape(batchSizes.begin(), batchSizes.end());
+      shape.append(T.getShape().begin(), T.getShape().end());
+      Ty = T.clone(shape);
+    }
+    mlir::NamedAttrList attrs;
+    for (auto attr : src->getAttrs()) {
+      auto eattr = cast<DenseI64ArrayAttr>(attr.getValue());
+      SmallVector<int64_t> shape;
+      for (size_t i = 0; i < batchSizes.size(); i++)
+        shape.push_back(i);
+      for (auto val : eattr.asArrayRef())
+        shape.push_back(val + batchSizes.size());
+      attr.setValue(DenseI64ArrayAttr::get(src->getContext(), shape));
+      attrs.append(attr);
+    }
+    auto cop = mlir::Operation::create(
+        src->getLoc(), src->getName(), resultTypes, {}, std::move(attrs),
+        OpaqueProperties(nullptr), mlir::BlockRange(), 0);
+    return cop;
+  }
+};
+
+struct ADDataFlowSortOp
+    : public ADDataFlowOpInterface::ExternalModel<ADDataFlowSortOp, SortOp> {
+
+  SmallVector<Value> getPotentialIncomingValuesRes(Operation *op,
+                                                   mlir::OpResult v) const {
+    auto srt = cast<SortOp>(op);
+    return {srt.getInputs()[v.getResultNumber()]};
+  }
+  SmallVector<Value>
+  getPotentialIncomingValuesArg(Operation *op, mlir::BlockArgument v) const {
+    auto srt = cast<SortOp>(op);
+    size_t num = v.getArgNumber() / 2;
+    return {srt.getInputs()[num]};
+  }
+  // The return of region is a comparison, non differentiable, and non flowing
+  SmallVector<Value> getPotentialTerminatorUsers(Operation *op, Operation *term,
+                                                 Value v) const {
+    return {};
+  }
+};
+
+struct RegionBranchCaseOp
+    : public RegionBranchOpInterface::ExternalModel<RegionBranchCaseOp,
+                                                    CaseOp> {
+
+  void
+  getEntrySuccessorRegions(Operation *op, ArrayRef<Attribute> operands,
+                           SmallVectorImpl<RegionSuccessor> &successors) const {
+    for (auto &reg : op->getRegions())
+      successors.push_back(RegionSuccessor(&reg));
+  }
+
+  mlir::OperandRange getEntrySuccessorOperands(Operation *op,
+                                               RegionBranchPoint bp) const {
+    auto end = op->operand_end();
+    return ::mlir::OperandRange(end, end);
+  }
+
+  void
+  getRegionInvocationBounds(Operation *op, ArrayRef<Attribute> operands,
+                            SmallVectorImpl<InvocationBounds> &bounds) const {
+    bounds.append(op->getNumRegions(), InvocationBounds(/*lb=*/0, /*ub=*/1));
+  }
+
+  bool areTypesCompatible(Operation *op, Type lhs, Type rhs) const {
+    return lhs == rhs;
+  }
+
+  void getSuccessorRegions(Operation *op, RegionBranchPoint point,
+                           SmallVectorImpl<RegionSuccessor> &regions) const {
+    // The `then` and the `else` region branch back to the parent operation.
+    if (!point.isParent()) {
+      regions.push_back(RegionSuccessor(op->getResults()));
+      return;
+    }
+
+    for (auto &reg : op->getRegions())
+      regions.push_back(RegionSuccessor(&reg));
+  }
+};
+
+struct ScatterActivity
+    : public ActivityOpInterface::ExternalModel<ScatterActivity, ScatterOp> {
+  bool isInactive(Operation *op) const { return false; }
+  bool isArgInactive(Operation *op, size_t idx) const {
+    auto scat = cast<ScatterOp>(op);
+    if (idx >= scat.getInputs().getBeginOperandIndex() &&
+        idx < scat.getInputs().getBeginOperandIndex() + scat.getInputs().size())
+      return false;
+    if (idx >= scat.getUpdates().getBeginOperandIndex() &&
+        idx <
+            scat.getUpdates().getBeginOperandIndex() + scat.getUpdates().size())
+      return false;
+    return true;
+  }
+};
+
+struct ADDataFlowWhileOp
+    : public ADDataFlowOpInterface::ExternalModel<ADDataFlowWhileOp, WhileOp> {
+
+  SmallVector<Value> getPotentialIncomingValuesRes(Operation *op,
+                                                   mlir::OpResult v) const {
+    auto srt = cast<WhileOp>(op);
+    auto resvals = cast<ReturnOp>(srt.getBody().front().back());
+    return {
+        srt.getOperand()[v.getResultNumber()],
+        resvals.getOperands()[v.getResultNumber()],
+    };
+  }
+
+  SmallVector<Value>
+  getPotentialIncomingValuesArg(Operation *op, mlir::BlockArgument v) const {
+    auto srt = cast<WhileOp>(op);
+    auto resvals = cast<ReturnOp>(srt.getBody().front().back());
+    return {
+        srt.getOperand()[v.getArgNumber()],
+        resvals.getOperands()[v.getArgNumber()],
+    };
+  }
+  SmallVector<Value> getPotentialTerminatorUsers(Operation *op, Operation *term,
+                                                 Value v) const {
+    auto srt = cast<WhileOp>(op);
+    if (&srt.getCond() == term->getParentRegion())
+      return {};
+    SmallVector<Value> sv;
+    for (auto &&[res, arg] : llvm::zip(srt.getResults(), term->getOperands())) {
+      if (arg == v)
+        sv.push_back(res);
+    }
+    return sv;
+  }
+};
+
+struct ADDataFlowReduceOp
+    : public ADDataFlowOpInterface::ExternalModel<ADDataFlowReduceOp,
+                                                  ReduceOp> {
+
+  SmallVector<Value> getPotentialIncomingValuesRes(Operation *op,
+                                                   mlir::OpResult v) const {
+    auto srt = cast<ReduceOp>(op);
+    auto resvals = cast<ReturnOp>(op->getRegion(0).front().back());
+    return {
+        srt.getInputs()[v.getResultNumber()],
+        srt.getInitValues()[v.getResultNumber()],
+        resvals.getOperands()[v.getResultNumber()],
+    };
+  }
+  SmallVector<Value>
+  getPotentialIncomingValuesArg(Operation *op, mlir::BlockArgument v) const {
+    auto srt = cast<ReduceOp>(op);
+    auto resvals = cast<ReturnOp>(op->getRegion(0).front().back());
+    if (v.getArgNumber() < srt.getInitValues().size())
+      return {srt.getInitValues()[v.getArgNumber()],
+              resvals.getOperands()[v.getArgNumber()]};
+    else
+      return {
+          srt.getInputs()[v.getArgNumber() - srt.getInitValues().size()],
+          resvals.getOperands()[v.getArgNumber() - srt.getInitValues().size()]};
+  }
+  SmallVector<Value> getPotentialTerminatorUsers(Operation *op, Operation *term,
+                                                 Value v) const {
+    auto srt = cast<ReduceOp>(op);
+    SmallVector<Value> sv;
+    for (size_t i = 0; i < srt.getInputs().size(); i++) {
+      if (term->getOperands()[i] == v) {
+        sv.push_back(srt.getResults()[i]);
+        sv.push_back(op->getRegion(0)
+                         .front()
+                         .getArguments()[srt.getInitValues().size() + i]);
+      }
+    }
+    return sv;
+  }
+};
+
+struct ADDataFlowScatterOp
+    : public ADDataFlowOpInterface::ExternalModel<ADDataFlowScatterOp,
+                                                  ScatterOp> {
+
+  SmallVector<Value> getPotentialIncomingValuesRes(Operation *op,
+                                                   mlir::OpResult v) const {
+    auto srt = cast<ScatterOp>(op);
+    auto resvals = cast<ReturnOp>(op->getRegion(0).front().back());
+    return {
+        srt.getInputs()[v.getResultNumber()],
+        resvals.getOperands()[v.getResultNumber()],
+    };
+  }
+  SmallVector<Value>
+  getPotentialIncomingValuesArg(Operation *op, mlir::BlockArgument v) const {
+    auto srt = cast<ScatterOp>(op);
+    auto resvals = cast<ReturnOp>(op->getRegion(0).front().back());
+    if (v.getArgNumber() < srt.getInputs().size())
+      return {resvals.getOperands()[v.getArgNumber()],
+              srt.getInputs()[v.getArgNumber()]};
+    else
+      return {srt.getUpdates()[v.getArgNumber() - srt.getInputs().size()]};
+  }
+  SmallVector<Value> getPotentialTerminatorUsers(Operation *op, Operation *term,
+                                                 Value v) const {
+    auto srt = cast<ScatterOp>(op);
+    SmallVector<Value> sv;
+    for (size_t i = 0; i < srt.getInputs().size(); i++) {
+      if (term->getOperands()[i] == v) {
+        sv.push_back(srt.getResults()[i]);
+      }
+    }
+    return sv;
+  }
+};
+
+class AutoDiffScatter
+    : public AutoDiffOpInterface::ExternalModel<AutoDiffScatter, ScatterOp> {
+public:
+  LogicalResult createForwardModeTangent(Operation *op, OpBuilder &builder,
+                                         MGradientUtils *gutils) const {
+    auto scat = cast<ScatterOp>(op);
+    SmallVector<int> inds;
+
+    for (auto idx = scat.getInputs().getBeginOperandIndex();
+         idx <
+         scat.getInputs().getBeginOperandIndex() + scat.getInputs().size();
+         idx++)
+      inds.push_back((int)idx);
+
+    for (auto idx = scat.getUpdates().getBeginOperandIndex();
+         idx <
+         scat.getUpdates().getBeginOperandIndex() + scat.getUpdates().size();
+         idx++)
+      inds.push_back((int)idx);
+
+    SmallVector<Type> ResultTypes;
+    SmallVector<Value> Inputs;
+    SmallVector<Value> Updates;
+    for (auto &&[res, inp, up] :
+         llvm::zip(op->getResults(), scat.getInputs(), scat.getUpdates())) {
+      ResultTypes.push_back(res.getType());
+      Inputs.push_back(gutils->getNewFromOriginal(inp));
+      Updates.push_back(gutils->getNewFromOriginal(up));
+      if (!gutils->isConstantValue(res)) {
+        ResultTypes.push_back(
+            cast<AutoDiffTypeInterface>(res.getType()).getShadowType());
+        if (!gutils->isConstantValue(inp)) {
+          Inputs.push_back(gutils->invertPointerM(inp, builder));
+        } else {
+          Inputs.push_back(cast<AutoDiffTypeInterface>(inp.getType())
+                               .createNullValue(builder, op->getLoc()));
+        }
+        if (!gutils->isConstantValue(up)) {
+          Updates.push_back(gutils->invertPointerM(up, builder));
+        } else {
+          Updates.push_back(cast<AutoDiffTypeInterface>(up.getType())
+                                .createNullValue(builder, op->getLoc()));
+        }
+      }
+    }
+
+    auto replacement = builder.create<ScatterOp>(
+        scat.getLoc(), ResultTypes, Inputs,
+        gutils->getNewFromOriginal(scat.getScatterIndices()), Updates,
+        scat.getScatterDimensionNumbersAttr(), scat.getIndicesAreSortedAttr(),
+        scat.getUniqueIndicesAttr());
+    auto newOp = gutils->getNewFromOriginal(op);
+    for (auto &&[region, replacementRegion] :
+         llvm::zip(newOp->getRegions(), replacement->getRegions())) {
+      replacementRegion.takeBody(region);
+    }
+
+    auto nb = gutils->getNewFromOriginal(&scat->getRegion(0).front());
+    // Rewrite block arguments to match the shadowing
+    size_t curidx = 0;
+    for (int j = 0; j < 2; j++) {
+      for (size_t i = 0; i < scat.getInputs().size(); i++) {
+        auto inp = scat.getInputs()[i];
+        auto up = scat.getUpdates()[i];
+        // primal
+        curidx++;
+        if (gutils->isConstantValue(inp) && gutils->isConstantValue(up)) {
+          continue;
+        }
+        auto ba = scat->getRegion(0)
+                      .front()
+                      .getArguments()[i + j * scat.getInputs().size()];
+        if (gutils->isConstantValue(ba)) {
+          nb->insertArgument(
+              curidx, cast<AutoDiffTypeInterface>(ba.getType()).getShadowType(),
+              scat.getLoc());
+        }
+        // shadow
+        curidx++;
+      }
+    }
+
+    // Inject the mapping for the new results into GradientUtil's shadow
+    // table.
+    SmallVector<Value> reps;
+    size_t idx = 0;
+    for (Value r : op->getResults()) {
+      // TODO only if used
+      reps.push_back(replacement->getResult(idx));
+      idx++;
+      if (!gutils->isConstantValue(r)) {
+        auto inverted = gutils->invertedPointers.lookupOrNull(r);
+        assert(inverted);
+        gutils->invertedPointers.map(r, replacement->getResult(idx));
+        inverted.replaceAllUsesWith(replacement->getResult(idx));
+        gutils->erase(inverted.getDefiningOp());
+        idx++;
+      }
+    }
+
+    // Differentiate body.
+    for (auto &origRegion : op->getRegions()) {
+      for (auto &origBlock : origRegion) {
+        for (Operation &o : origBlock) {
+          if (failed(gutils->visitChild(&o))) {
+            return failure();
+          }
+        }
+      }
+    }
+
+    // Replace all uses of original results
+    gutils->replaceOrigOpWith(op, reps);
+    gutils->erase(newOp);
+    gutils->originalToNewFnOps[op] = replacement;
+    return success();
+  }
+};
+
+class AutoDiffHLOReturn
+    : public AutoDiffOpInterface::ExternalModel<AutoDiffHLOReturn, ReturnOp> {
+public:
+  LogicalResult createForwardModeTangent(Operation *origTerminator,
+                                         OpBuilder &builder,
+                                         MGradientUtils *gutils) const {
+    auto parentOp = origTerminator->getParentOp();
+
+    llvm::SmallDenseSet<unsigned> operandsToPrimal;
+    llvm::SmallDenseSet<unsigned> operandsToShadow;
+    for (OpOperand &operand : origTerminator->getOpOperands()) {
+      operandsToPrimal.insert(operand.getOperandNumber());
+    }
+    if (auto wop = dyn_cast<WhileOp>(parentOp)) {
+      if (origTerminator->getParentRegion() ==
+          &cast<WhileOp>(parentOp).getCond()) {
+      } else {
+        llvm::SmallDenseSet<unsigned> resultPositionsToShadow;
+
+        for (auto &&[res, arg] :
+             llvm::zip(wop.getResults(), wop.getBody().front().getArguments()))
+          if (!gutils->isConstantValue(res) || !gutils->isConstantValue(arg)) {
+            operandsToShadow.insert(res.getResultNumber());
+          }
+      }
+    } else if (isa<FunctionOpInterface>(parentOp)) {
+      operandsToPrimal.clear();
+      for (OpOperand &operand : origTerminator->getOpOperands()) {
+        auto idx = operand.getOperandNumber();
+        if (gutils->returnPrimals[idx])
+          operandsToPrimal.insert(idx);
+        if (gutils->returnShadows[idx])
+          operandsToShadow.insert(idx);
+      }
+    } else {
+      assert(parentOp->getNumResults() == origTerminator->getNumOperands());
+      for (auto res : parentOp->getResults()) {
+        if (!gutils->isConstantValue(res))
+          operandsToShadow.insert(res.getResultNumber());
+      }
+    }
+
+    SmallVector<Value> newOperands;
+    newOperands.reserve(operandsToPrimal.size() + operandsToShadow.size());
+    for (OpOperand &operand : origTerminator->getOpOperands()) {
+      if (operandsToPrimal.contains(operand.getOperandNumber()))
+        newOperands.push_back(gutils->getNewFromOriginal(operand.get()));
+      if (operandsToShadow.contains(operand.getOperandNumber())) {
+        if (!gutils->isConstantValue(operand.get())) {
+          newOperands.push_back(gutils->invertPointerM(operand.get(), builder));
+        } else {
+          Type retTy = operand.get()
+                           .getType()
+                           .cast<AutoDiffTypeInterface>()
+                           .getShadowType();
+          newOperands.push_back(
+              retTy.cast<AutoDiffTypeInterface>().createNullValue(
+                  builder, origTerminator->getLoc()));
+        }
+      }
+    }
+
+    // Assuming shadows following the originals are fine.
+    // TODO: consider extending to have a ShadowableTerminatorOpInterface
+    Operation *replTerminator = gutils->getNewFromOriginal(origTerminator);
+    replTerminator->setOperands(newOperands);
+    return success();
+  }
+};
+
+class AutoDiffSort
+    : public AutoDiffOpInterface::ExternalModel<AutoDiffSort, SortOp> {
+public:
+  LogicalResult createForwardModeTangent(Operation *op, OpBuilder &builder,
+                                         MGradientUtils *gutils) const {
+
+    // TODO: we may need to record, for every successor, which of its inputs
+    // need a shadow to recreate the body correctly.
+    llvm::SmallDenseSet<unsigned> operandPositionsToShadow;
+    llvm::SmallDenseSet<unsigned> resultPositionsToShadow;
+
+    for (auto res : op->getResults())
+      if (!gutils->isConstantValue(res)) {
+        operandPositionsToShadow.insert(res.getResultNumber());
+        resultPositionsToShadow.insert(res.getResultNumber());
+      }
+
+    return mlir::enzyme::detail::controlFlowForwardHandler(
+        op, builder, gutils, operandPositionsToShadow, resultPositionsToShadow);
+  }
+};
+
 } // namespace
 
 void mlir::enzyme::registerStableHLODialectAutoDiffInterface(
@@ -506,12 +989,30 @@ void mlir::enzyme::registerStableHLODialectAutoDiffInterface(
   registry.addExtension(
       +[](MLIRContext *context, stablehlo::StablehloDialect *) {
         registerInterfaces(context);
+
+        // SortOp::attachInterface<AutoDiffSort>(*context);
+
+        WhileOp::attachInterface<ADDataFlowWhileOp>(*context);
+        SortOp::attachInterface<ADDataFlowSortOp>(*context);
+        ScatterOp::attachInterface<ADDataFlowScatterOp>(*context);
+        ReduceOp::attachInterface<ADDataFlowReduceOp>(*context);
+
+        CaseOp::attachInterface<RegionBranchCaseOp>(*context);
+
+        ScatterOp::attachInterface<ScatterActivity>(*context);
+        ScatterOp::attachInterface<AutoDiffScatter>(*context);
+
+        ReturnOp::attachInterface<AutoDiffHLOReturn>(*context);
+
         ReduceOp::attachInterface<AutoDiffReduceFwd<ReduceOp>>(*context);
+        WhileOp::attachInterface<AutoDiffWhileFwd>(*context);
         ReduceOp::attachInterface<AutoDiffReduceCF<ReduceOp>>(*context);
+        WhileOp::attachInterface<AutoDiffReduceCF<WhileOp>>(*context);
         BroadcastInDimOp::attachInterface<AutoDiffBroadcastInDimRev>(*context);
         SliceOp::attachInterface<AutoDiffSliceRev>(*context);
         ReduceOp::attachInterface<AutoDiffReduceRev>(*context);
         ConcatenateOp::attachInterface<AutoDiffConcatenateRev>(*context);
         ConstantOp::attachInterface<SHLOConstantOpBatchInterface>(*context);
+        TransposeOp::attachInterface<SHLOTransposeOpBatchInterface>(*context);
       });
 }
diff --git a/src/enzyme_ad/jax/Implementations/StableHLODerivatives.td b/src/enzyme_ad/jax/Implementations/StableHLODerivatives.td
index 86c30c49..38029ec8 100644
--- a/src/enzyme_ad/jax/Implementations/StableHLODerivatives.td
+++ b/src/enzyme_ad/jax/Implementations/StableHLODerivatives.td
@@ -34,3 +34,14 @@ def : HLODerivative<"UnaryEinsumOp", (Op $x),
                     ],
                     (UnaryEinsum (ResultTypes), (Shadow $x), (EinsumConfig))
                   >;
+
+def : ControlFlowOp<"stablehlo", "CaseOp", [{
+  Operation *createWithShadows(Operation *op, OpBuilder &builder,
+                               MGradientUtils *gutils, Operation *original,
+                               ValueRange remappedOperands,
+                               TypeRange rettys) const {
+    return builder.create<stablehlo::CaseOp>(original->getLoc(), rettys,
+                                        remappedOperands, original->getAttrs(), cast<stablehlo::CaseOp>(op).getBranches().size());
+  }
+}]>;
+
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
index d6126647..3f2a1177 100644
--- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
+++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -3011,6 +3011,24 @@ struct ConvertSimplify : public OpRewritePattern<mlir::stablehlo::ConvertOp> {
 struct SliceSimplify : public OpRewritePattern<mlir::stablehlo::SliceOp> {
   using OpRewritePattern<mlir::stablehlo::SliceOp>::OpRewritePattern;
 
+  static size_t getDenseElementBitWidth(Type eltType) {
+    // Align the width for complex to 8 to make storage and interpretation
+    // easier.
+    if (ComplexType comp = llvm::dyn_cast<ComplexType>(eltType))
+      return llvm::alignTo<8>(getDenseElementBitWidth(comp.getElementType())) *
+             2;
+    if (eltType.isIndex())
+      return IndexType::kInternalStorageBitWidth;
+    return eltType.getIntOrFloatBitWidth();
+  }
+
+  static size_t getDenseElementStorageWidth(size_t origWidth) {
+    return origWidth == 1 ? origWidth : llvm::alignTo<8>(origWidth);
+  }
+  static size_t getDenseElementStorageWidth(Type elementType) {
+    return getDenseElementStorageWidth(getDenseElementBitWidth(elementType));
+  }
+
   LogicalResult matchAndRewrite(mlir::stablehlo::SliceOp op,
                                 PatternRewriter &rewriter) const final {
     DenseElementsAttr inp;
@@ -3020,12 +3038,41 @@ struct SliceSimplify : public OpRewritePattern<mlir::stablehlo::SliceOp> {
       if (inp.isSplat()) {
         out = inp.resizeSplat(op.getType());
       } else {
-        auto ten = mlir::stablehlo::constantOp(inp);
-        out = fromTensor(mlir::stablehlo::sliceOp(
-            ten, stablehlo::Sizes(op.getStartIndices()),
-            stablehlo::Sizes(op.getStrides()), op.getType()));
+        bool contiguous = true;
+        size_t offset = 0;
+        auto inshape = op.getOperand().getType().getShape();
+        auto outshape = op.getType().getShape();
+        size_t total = 1;
+        for (int i = 0; i < inshape.size(); i++) {
+          if (op.getStrides()[i] != 1) {
+            contiguous = false;
+          }
+          auto start = op.getStartIndices()[i];
+          auto lim = op.getLimitIndices()[i];
+          if (start != 0 || lim != inshape[i]) {
+            if (offset != 0) {
+              contiguous = false;
+            }
+          }
+          offset *= inshape[i];
+          offset += start;
+          total *= outshape[i];
+        }
+        auto elementType = op.getOperand().getType().getElementType();
+        auto bw = getDenseElementStorageWidth(elementType);
+        if (contiguous && bw != 1) {
+          const char *elementPtr = inp.getRawData().data() + (bw / 8) * offset;
+
+          auto values = ArrayRef((char *)elementPtr, (bw / 8) * total);
+          out =
+              DenseIntOrFPElementsAttr::getFromRawBuffer(op.getType(), values);
+        } else {
+          auto ten = mlir::stablehlo::constantOp(inp);
+          out = fromTensor(mlir::stablehlo::sliceOp(
+              ten, stablehlo::Sizes(op.getStartIndices()),
+              stablehlo::Sizes(op.getStrides()), op.getType()));
+        }
       }
-
       rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(), out);
       return success();
     }
@@ -3073,25 +3120,6 @@ struct BroadcastInDimSimplify
   }
 };
 
-struct ReshapeSimplify : public OpRewritePattern<mlir::stablehlo::ReshapeOp> {
-  using OpRewritePattern<mlir::stablehlo::ReshapeOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(mlir::stablehlo::ReshapeOp op,
-                                PatternRewriter &rewriter) const final {
-    DenseElementsAttr inp;
-    matchPattern(op->getOperand(0), m_Constant(&inp));
-    if (inp) {
-      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
-          op, op.getType(),
-          inp.isSplat() ? inp.resizeSplat(op.getType())
-                        : inp.reshape(op.getType()));
-      return success();
-    }
-
-    return failure();
-  }
-};
-
 struct DotGeneralSimplify
     : public OpRewritePattern<mlir::stablehlo::DotGeneralOp> {
   using OpRewritePattern<mlir::stablehlo::DotGeneralOp>::OpRewritePattern;
@@ -6170,21 +6198,12 @@ struct ReshapeOpCanon final : OpRewritePattern<mlir::stablehlo::ReshapeOp> {
     }
 
     // Fold reshape of a constant.
-    ElementsAttr cstAttr;
+    DenseElementsAttr cstAttr;
     if (!matchPattern(op.getOperand(), m_Constant(&cstAttr)))
       return failure();
 
-    if (auto splat = cstAttr.dyn_cast<SplatElementsAttr>()) {
-      rewriter.replaceOpWithNewOp<mlir::stablehlo::ConstantOp>(
-          op, SplatElementsAttr::get(op.getType(),
-                                     splat.getSplatValue<Attribute>()));
-      return success();
-    }
-
-    auto elements =
-        llvm::to_vector_of<Attribute>(cstAttr.getValues<Attribute>());
     rewriter.replaceOpWithNewOp<mlir::stablehlo::ConstantOp>(
-        op, DenseElementsAttr::get(op.getType(), elements));
+        op, cstAttr.reshape(op.getType()));
     return success();
   }
 };
@@ -6394,11 +6413,11 @@ struct EnzymeHLOOptPass : public EnzymeHLOOptPassBase<EnzymeHLOOptPass> {
              PowSimplify, SqrtSimplify, CosSimplify, SinSimplify, NoopSlice,
              SliceSlice, PadSimplify, ShiftRightLogicalSimplify,
              NegativePadToSlice, TanhSimplify, ExpSimplify, SliceSimplify,
-             ConvertSimplify, ReshapeSimplify, TransposeSimplify,
-             DotGeneralSimplify, DynamicSliceToStatic, DynamicUpdateSliceElim,
-             ReduceToReshape, BroadcastToReshape, GatherSimplify,
-             ReshapeEmptyBroadcast, BroadcastReshape, ConstPropThroughBarrier>(
-            context, PatternBenefit(65000));
+             ConvertSimplify, TransposeSimplify, DotGeneralSimplify,
+             DynamicSliceToStatic, DynamicUpdateSliceElim, ReduceToReshape,
+             BroadcastToReshape, GatherSimplify, ReshapeEmptyBroadcast,
+             BroadcastReshape, ConstPropThroughBarrier>(context,
+                                                        PatternBenefit(65000));
 
     patterns.add<IotaSimplify, BroadcastInDimSimplify>(
         max_constant_expansion, context, PatternBenefit(65000));
diff --git a/src/enzyme_ad/jax/Passes/EqualitySaturation.cpp b/src/enzyme_ad/jax/Passes/EqualitySaturation.cpp
index 59017903..c0e16000 100644
--- a/src/enzyme_ad/jax/Passes/EqualitySaturation.cpp
+++ b/src/enzyme_ad/jax/Passes/EqualitySaturation.cpp
@@ -15,19 +15,29 @@
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
+
 #include "src/enzyme_ad/jax/deps/include/ReactantExtra.h"
 #include "stablehlo/dialect/StablehloOps.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
 #include "xla/pjrt/cpu/cpu_client.h"
 #include "xla/pjrt/gpu/se_gpu_pjrt_client.h"
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/pjrt/status_casters.h"
 #include "xla/python/ifrt/executable.h"
 #include "xla/python/pjrt_ifrt/xla_compiler.h"
+#include "xla/service/gpu/gpu_latency_hiding_scheduler.h"
+#include "xla/service/gpu/model/gpu_hlo_cost_analysis.h"
+#include "xla/service/gpu/model/gpu_performance_model.h"
 #include "xla/service/platform_util.h"
+#include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/translate/mhlo_to_hlo/mlir_hlo_to_hlo.h"
 
 #include "cxxbridge/deps/tensat/src/input.rs.h"
 #include "rust/cxx.h"
@@ -41,6 +51,9 @@
 #include <string>
 #include <thread>
 
+#undef NDEBUG
+#include <cassert>
+
 #define DEBUG_TYPE "enzyme"
 
 using namespace mlir;
@@ -111,6 +124,29 @@ EqsatPlatform getPlatform() {
   }
 }
 
+stream_executor::DeviceDescription
+RTXA6000DeviceInfo(tensorflow::se::GpuComputeCapability cc =
+                       tensorflow::se::CudaComputeCapability(8, 9)) {
+  stream_executor::DeviceDescription b;
+  b.set_gpu_compute_capability(cc);
+  b.set_threads_per_block_limit(1024);
+  b.set_threads_per_warp(32);
+  b.set_shared_memory_per_block(48 * 1024);
+  b.set_shared_memory_per_block_optin(99 * 1024);
+  b.set_shared_memory_per_core(100 * 1024);
+  b.set_threads_per_core_limit(1536);
+  b.set_core_count(84);
+  b.set_fpus_per_core(128);
+  b.set_block_dim_limit_x(2'147'483'647);
+  b.set_block_dim_limit_y(65535);
+  b.set_block_dim_limit_z(65535);
+  b.set_memory_bandwidth(768'096'000'000);
+  b.set_l2_cache_size(6 * 1024 * 1024);
+  b.set_clock_rate_ghz(1.410);
+  b.set_device_memory_size(51'050'250'240);
+  return b;
+}
+
 /**
  * Ops to not measure cost for
  */
@@ -147,11 +183,17 @@ bool isBlackboxed(Operation *op) {
 class OperationTimer {
 public:
   /**
-   * Measure cost of operation (execution time in microseconds) by running it
+   * Get cost of operation. This depends on the platform:
+   * - if CPU, then we measure the execution time in microseconds by running it
    * many times and measuring the time taken.
+   * - if GPU, we use XLA's analytical cost model to get the expected execution
+   * time in nanoseconds.
    */
   static uint64_t getCost(Operation *op, unsigned warmup,
                           unsigned repetitions) {
+    // TODO: refactor this/separate out into two functions? so that
+    // warmup/repetitions don't need to be passed for GPU.
+
     if (!logsInitialized) {
       InitializeLogs();
       logsInitialized = true;
@@ -169,6 +211,7 @@ class OperationTimer {
     }
 
     auto context = OperationTimer::getContext();
+    int cost = -1;
 
     // For some reason, not cloning the op here leads to a segfault. Maybe
     // prepareExecutable/ClientCompile consumes the op?
@@ -177,72 +220,88 @@ class OperationTimer {
     ModuleOp wrapperModule =
         createModuleFromOperation(context, opForMeasurement);
 
-    xla::PjRtClient *client = nullptr;
+    switch (getPlatform()) {
+    case CPU: {
+      xla::PjRtClient *client = MakeCPUClient(0, 1, 1);
+      auto executable = prepareExecutable(client, wrapperModule);
 
-    auto platform = getPlatform();
+      unsigned numArgs = op->getNumOperands();
+      unsigned numResults = op->getNumResults();
+      uint8_t futures = 0;
 
-    switch (platform) {
-    case CPU:
-      client = MakeCPUClient(0, 1, 1);
-      break;
-    case GPU:
-      // https://github.com/EnzymeAD/Reactant.jl/blob/65060404e19cd5a56a51e4fb2b252380477632b0/src/XLA.jl#L56
-      const char *error = "";
-      // TODO: is this correct?
-      client = MakeGPUClient(0, 1, nullptr, 0, "gpu", &error);
-      if (std::string(error) != "") {
-        auto error_string =
-            "Error while creating GPU client: " + std::string(error);
-        throw std::invalid_argument(error_string);
-      }
-      break;
-    }
+      xla::PjRtBuffer *args[numArgs];
+      uint8_t isArgDonatable[numArgs];
 
-    auto executable = prepareExecutable(client, wrapperModule);
+      int numRuns = warmup + repetitions;
+      xla::PjRtBuffer *res[numRuns * numResults];
 
-    unsigned numArgs = op->getNumOperands();
-    unsigned numResults = op->getNumResults();
-    uint8_t futures = 0;
+      for (int i = 0; i < numArgs; i++) {
+        args[i] = getRandomInput(client, op->getOperand(i).getType());
+        isArgDonatable[i] = false;
+      }
 
-    xla::PjRtBuffer *args[numArgs];
-    uint8_t isArgDonatable[numArgs];
+      std::vector<uint64_t> durations(numRuns);
 
-    int numRuns = warmup + repetitions;
-    xla::PjRtBuffer *res[numRuns * numResults];
+      for (unsigned i = 0; i < warmup + repetitions; i++) {
+        durations[i] =
+            XLAExecute(executable, numArgs, args, isArgDonatable, numResults,
+                       res + i * numResults, &futures, nullptr);
+      }
 
-    for (int i = 0; i < numArgs; i++) {
-      args[i] = getRandomInput(client, op->getOperand(i).getType());
-      isArgDonatable[i] = false;
-    }
+      // TODO: This means there's no point in warmup anymore, since we're now
+      // taking individual measurements. Maybe we get rid of the parameter or do
+      // something more sophisticated
+      cost = *std::min_element(durations.begin(), durations.end());
+      assert(!futures);
 
-    std::vector<uint64_t> durations(numRuns);
+      // Cleanup
+      for (int i = 0; i < numRuns * numResults; i++) {
+        PjRtBufferFree(res[i]);
+      }
 
-    for (unsigned i = 0; i < warmup + repetitions; i++) {
-      durations[i] =
-          XLAExecute(executable, numArgs, args, isArgDonatable, numResults,
-                     res + i * numResults, &futures, nullptr);
+      FreeClient(executable->client());
+      ExecutableFree(executable);
+      break;
+    }
+    case GPU: {
+      std::unique_ptr<xla::HloModule> hloModule =
+          wrapperModuleToHloModule(wrapperModule);
+
+      auto deviceDescription = getDeviceDescription();
+
+      xla::HloCostAnalysis::ShapeSizeFunction shapeSizeFunction =
+          [](const xla::Shape &shape) {
+            return xla::gpu::GetSizeOfShape(shape, 4);
+          };
+      xla::gpu::GpuHloCostAnalysis costAnalysis(
+          xla::gpu::GpuHloCostAnalysis::Options{
+              shapeSizeFunction, {}, {}, true},
+          *deviceDescription);
+
+      assert(hloModule->computation_count() == 1);
+      for (auto c : hloModule->computations()) {
+        c->Accept(&costAnalysis);
+        // The op we are measuring should always be the return value, which is
+        // at the root.
+        auto op = c->root_instruction();
+
+        auto runtime =
+            xla::gpu::GpuPerformanceModel::EstimateRunTimeForInstruction(
+                op, *deviceDescription, &costAnalysis,
+                xla::gpu::GpuPerformanceModelOptions::ForModule(
+                    op->GetModule()));
+        cost = absl::ToInt64Nanoseconds(runtime.exec_time);
+      }
+      break;
+    }
     }
 
-    assert(!futures);
-
-    // TODO: This means there's no point in warmup anymore, since we're now
-    // taking individual measurements. Maybe we get rid of the parameter or do
-    // something more sophisticated
-    uint64_t duration = *std::min_element(durations.begin(), durations.end());
-
+    assert(cost >= 0);
     auto indexOp = op->clone();
-    runtimeCache.try_emplace(indexOp, duration);
-
-    // Cleanup
-    for (int i = 0; i < numRuns * numResults; i++) {
-      PjRtBufferFree(res[i]);
-    }
+    runtimeCache.try_emplace(indexOp, cost);
 
-    FreeClient(executable->client());
-    ExecutableFree(executable);
     wrapperModule.erase();
-
-    return duration;
+    return cost;
   }
 
   static Operation *getDummyOp(OpBuilder &builder, Type type) {
@@ -372,6 +431,42 @@ class OperationTimer {
 
     return buffer;
   }
+
+  /**
+   * Create XLA internal HloModule for the analytical cost model
+   */
+  static std::unique_ptr<xla::HloModule>
+  wrapperModuleToHloModule(ModuleOp &wrapperModule) {
+    auto context = wrapperModule.getContext();
+    PassManager pm(context);
+    pm.addPass(mlir::mhlo::createStablehloLegalizeToHloPass());
+    pm.run(wrapperModule);
+
+    MlirToHloConversionOptions options;
+    options.propagate_layouts = true;
+    options.return_tuple = false;
+
+    auto hloModule = ConvertMlirHloToHloModule(wrapperModule, options);
+    if (!hloModule.ok()) {
+      llvm::errs() << "Couldn't create hloModule: "
+                   << hloModule.status().message();
+      return nullptr;
+    } else {
+      return std::move(hloModule.value());
+    }
+  }
+
+  /**
+   * Get DeviceDescription for current device.
+   */
+  static std::unique_ptr<stream_executor::DeviceDescription>
+  getDeviceDescription() {
+    auto platform = xla::PlatformUtil::GetPlatform(
+                        (getPlatform() == EqsatPlatform::CPU ? "cpu" : "cuda"))
+                        .value();
+    // assume ordinal 0
+    return std::move(platform->DescriptionForDevice(0).value());
+  }
 };
 
 llvm::DenseMap<Operation *, uint64_t, OperationMapInfo>
@@ -446,6 +541,135 @@ std::vector<int64_t> dotGeneralShapeComputation(
 
   return shape;
 }
+/**
+ * Convert a vector of vectors into a DenseIntElementsAttr
+ */
+mlir::DenseIntElementsAttr
+matrixToDenseAttr(const std::vector<std::vector<int64_t>> &matrix,
+                  mlir::Builder &builder) {
+  std::vector<int64_t> flattenedResult;
+  std::vector<int64_t> shape;
+
+  if (matrix.empty()) {
+    std::cerr << "Error: Input matrix is empty." << std::endl;
+    return {};
+  }
+
+  for (const auto &row : matrix) {
+    if (row.empty()) {
+      std::cerr << "Error: Row in matrix is empty." << std::endl;
+      return {};
+    }
+
+    flattenedResult.insert(flattenedResult.end(), row.begin(), row.end());
+    shape.push_back(row.size());
+  }
+
+  if (shape.size() == 1) {
+    shape = {static_cast<long long>(flattenedResult.size())};
+  } else {
+    int64_t row_size = shape[0];
+    for (const auto &row : matrix) {
+      if (row.size() != row_size) {
+        std::cerr << "Error: Non-rectangular matrix detected." << std::endl;
+        return {};
+      }
+    }
+  }
+
+  int64_t expected_elements = 1;
+  for (auto dim : shape) {
+    expected_elements *= dim;
+  }
+
+  if (flattenedResult.size() != expected_elements) {
+    std::cerr << "Error: Mismatch between flattened result size and shape."
+              << " Expected " << expected_elements << " elements but got "
+              << flattenedResult.size() << "." << std::endl;
+    return {};
+  }
+  auto type = mlir::RankedTensorType::get(shape, builder.getIntegerType(64));
+  return mlir::DenseIntElementsAttr::get(type, flattenedResult);
+}
+
+// https://github.com/jax-ml/jax/blob/b8a066a90790a900d370812263dea51e4b262b43/jax/_src/lax/convolution.py#L351
+// Define ConvDimensionNumbers similar to JAX's implementation
+struct ConvDimensionNumbers {
+  std::vector<int64_t> lhs_spec; // (batch_dim, feature_dim, spatial_dims...)
+  std::vector<int64_t>
+      rhs_spec; // (out_feature_dim, in_feature_dim, spatial_dims...)
+  std::vector<int64_t> out_spec; // (batch_dim, feature_dim, spatial_dims...)
+};
+
+std::vector<int64_t> convolutionShapeComputation(
+    llvm::ArrayRef<int64_t> lhs_shape,    // Input shape
+    llvm::ArrayRef<int64_t> rhs_shape,    // Kernel shape
+    std::vector<int64_t> &window_strides, // Strides for each spatial dimension
+    mlir::DenseElementsAttr padding_attr, // Padding as DenseElementsAttr
+    std::vector<int64_t> &lhs_dilation,   // Dilation for input
+    std::vector<int64_t> &rhs_dilation,   // Dilation for kernel
+    ConvDimensionNumbers dimension_numbers,
+    int64_t feature_group_count, // Feature group count
+    int64_t batch_group_count    // Batch group count
+) {
+  // Ensure input shape has at least 2 dimensions
+  assert(lhs_shape.size() >= 2 &&
+         "Input shape must have at least 2 dimensions.");
+  assert(rhs_shape.size() >= 2 &&
+         "Kernel shape must have at least 2 dimensions.");
+
+  auto lhs_spec = dimension_numbers.lhs_spec;
+  auto rhs_spec = dimension_numbers.rhs_spec;
+  auto out_spec = dimension_numbers.out_spec;
+
+  std::vector<int64_t> lhs_spatial_dims(lhs_spec.begin() + 2, lhs_spec.end());
+  std::vector<int64_t> rhs_spatial_dims(rhs_spec.begin() + 2, rhs_spec.end());
+  std::vector<int64_t> out_spatial_dims(out_spec.begin() + 2, out_spec.end());
+
+  std::vector<std::pair<int64_t, int64_t>> padding;
+  auto padding_values = padding_attr.getValues<int64_t>();
+  auto it = padding_values.begin();
+  while (it != padding_values.end()) {
+    int64_t low = *it++;
+    int64_t high = *it++;
+    padding.emplace_back(low, high);
+  }
+
+  assert(padding.size() == lhs_spatial_dims.size() &&
+         "Padding size does not match number of spatial dimensions.");
+
+  int64_t output_batch_dim = lhs_shape[lhs_spec[0]] / batch_group_count;
+  int64_t output_feature_dim = rhs_shape[rhs_spec[0]] / feature_group_count;
+
+  std::vector<int64_t> output_shape(lhs_shape.size(), -1);
+  output_shape[out_spec[0]] = output_batch_dim;
+
+  // Compute spatial dimensions
+  for (size_t i = 0; i < lhs_spatial_dims.size(); ++i) {
+    int64_t lhs_dim = lhs_shape[lhs_spatial_dims[i]];
+    int64_t rhs_dim = rhs_shape[rhs_spatial_dims[i]];
+
+    // Apply dilation to input size
+    int64_t dilated_lhs = (lhs_dim - 1) * lhs_dilation[i] + 1;
+
+    // Apply padding
+    int64_t padded_lhs = dilated_lhs + padding[i].first + padding[i].second;
+
+    // Effective kernel size after dilation
+    int64_t dilated_rhs = (rhs_dim - 1) * rhs_dilation[i] + 1;
+
+    // Calculate output dimension
+    int64_t out_dim = (padded_lhs - dilated_rhs) / window_strides[i] + 1;
+    output_shape[out_spatial_dims[i]] = out_dim;
+  }
+  output_shape[out_spec[1]] = output_feature_dim;
+  std::cout << "Computed Output Shape: ";
+  for (int64_t dim : output_shape) {
+    std::cout << dim << " ";
+  }
+  std::cout << std::endl;
+  return output_shape;
+}
 
 /**
  * Get the correct start and limiting indices of a SliceOp from a SSplit0 or
@@ -456,18 +680,26 @@ std::vector<int64_t> dotGeneralShapeComputation(
 std::pair<std::vector<int64_t>, std::vector<int64_t>>
 getSliceIndicesFromSplit(tensat::Ops op, Value input, int axis, Value orig) {
   auto input_shape = getShape(input);
+  auto orig_shape = getShape(orig);
+
+  int input_rank = input_shape.size();
+  int orig_rank = orig_shape.size();
+
+  int axis1 = std::max(0, input_rank - orig_rank) + axis;
+
   std::vector<int64_t> start, limit;
+
   for (int i = 0; i < input_shape.size(); i++) {
-    if (i != axis) {
+    if (i != axis1) {
       start.push_back(0);
       limit.push_back(input_shape[i]);
     } else {
-      int slice_width = getShape(orig)[axis];
+      int slice_width = orig_shape[axis];
       if (op == tensat::Ops::SSplit0) {
         start.push_back(0);
         limit.push_back(slice_width);
       } else if (op == tensat::Ops::SSplit1) {
-        int input_width = input_shape[axis];
+        int input_width = input_shape[axis1];
         start.push_back(input_width - slice_width);
         limit.push_back(input_width);
       } else {
@@ -499,18 +731,25 @@ Type getReshapeTypeForMatchRank(Value input, Value ref) {
     for (int i = 0; i < initialRank; i++) {
       if (i < refRank)
         shape.push_back(initialShape[i]);
-      else
-        assert(initialShape[i] == 1);
+      else {
+        if (initialShape[i] != 1) {
+          throw std::invalid_argument(
+              "initial shape index " + std::to_string(i) +
+              " is not 1, found: " + std::to_string(initialShape[i]));
+        }
+      }
     }
   }
   return deriveOutputType(input, shape);
 }
 
-Operation *createStableHloOp(OpBuilder &builder, tensat::Ops op,
-                             SmallVector<Value> &operands,
-                             std::vector<std::vector<int64_t>> &other_vecs,
-                             std::vector<int64_t> &int_args,
-                             MLIRContext *context) {
+Operation *
+createStableHloOp(OpBuilder &builder, tensat::Ops op,
+                  SmallVector<Value> &operands,
+                  std::vector<std::vector<int64_t>> &other_vecs,
+                  std::vector<int64_t> &int_args,
+                  std::vector<std::vector<std::vector<int64_t>>> &matrix_args,
+                  MLIRContext *context) {
   Operation *mlirOp = nullptr;
 
   switch (op) {
@@ -567,6 +806,95 @@ Operation *createStableHloOp(OpBuilder &builder, tensat::Ops op,
                                                   newType, input);
     break;
   }
+  case tensat::Ops::ConvolutionOp: {
+    auto lhs = operands[0];
+    auto rhs = operands[1];
+
+    auto windowStrides = other_vecs[0];
+    auto padding = matrixToDenseAttr(matrix_args[0], builder);
+    auto lhsDilation = other_vecs[1];
+    auto rhsDilation = other_vecs[2];
+    auto windowReversal = other_vecs[3];
+    std::vector<uint8_t> windowReversalBool;
+    windowReversalBool.reserve(windowReversal.size());
+    for (int64_t val : windowReversal) {
+      windowReversalBool.push_back(static_cast<uint8_t>(val != 0));
+    }
+    // dimension numbers
+    auto inputBatchDimension = int_args[0];
+    auto inputFeatureDimension = int_args[1];
+    auto inputSpatialDimensions = other_vecs[4];
+    auto kernelInputFeatureDimension = int_args[2];
+    auto kernelOutputFeatureDimension = int_args[3];
+    auto kernelSpatialDimensions = other_vecs[5];
+    auto outputBatchDimension = int_args[4];
+    auto outputFeatureDimension = int_args[5];
+    auto outputSpatialDimensions = other_vecs[6];
+    auto convolutionDimensionNumbersAttr =
+        stablehlo::ConvDimensionNumbersAttr::get(
+            context, inputBatchDimension, inputFeatureDimension,
+            inputSpatialDimensions, kernelInputFeatureDimension,
+            kernelOutputFeatureDimension, kernelSpatialDimensions,
+            outputBatchDimension, outputFeatureDimension,
+            outputSpatialDimensions);
+
+    ConvDimensionNumbers convDims;
+    convDims.lhs_spec = {inputBatchDimension, inputFeatureDimension};
+    convDims.lhs_spec.insert(convDims.lhs_spec.end(),
+                             inputSpatialDimensions.begin(),
+                             inputSpatialDimensions.end());
+
+    convDims.rhs_spec = {kernelOutputFeatureDimension,
+                         kernelInputFeatureDimension};
+    convDims.rhs_spec.insert(convDims.rhs_spec.end(),
+                             kernelSpatialDimensions.begin(),
+                             kernelSpatialDimensions.end());
+
+    convDims.out_spec = {outputBatchDimension, outputFeatureDimension};
+    convDims.out_spec.insert(convDims.out_spec.end(),
+                             outputSpatialDimensions.begin(),
+                             outputSpatialDimensions.end());
+    auto featureGroupCount = int_args[6];
+    auto batchGroupCount = int_args[7];
+    auto precisionConfig = other_vecs[7];
+
+    std::vector<Attribute> precisionVec;
+    for (auto &precision : precisionConfig) {
+      switch (precision) {
+      case 0:
+        precisionVec.push_back(stablehlo::PrecisionAttr::get(
+            context, stablehlo::Precision::DEFAULT));
+        break;
+      case 1:
+        precisionVec.push_back(
+            stablehlo::PrecisionAttr::get(context, stablehlo::Precision::HIGH));
+        break;
+      case 2:
+        precisionVec.push_back(stablehlo::PrecisionAttr::get(
+            context, stablehlo::Precision::HIGHEST));
+        break;
+      }
+    }
+
+    auto shape = convolutionShapeComputation(
+        getShape(lhs), getShape(rhs), windowStrides, padding, lhsDilation,
+        rhsDilation, convDims, featureGroupCount, batchGroupCount);
+
+    auto newType = deriveOutputType(lhs, shape);
+    mlirOp = builder.create<stablehlo::ConvolutionOp>(
+        builder.getUnknownLoc(), newType, lhs, rhs,
+        mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(windowStrides)),
+        padding,
+        mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(lhsDilation)),
+        mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(rhsDilation)),
+        mlir::DenseBoolArrayAttr::get(
+            context, llvm::ArrayRef<bool>(reinterpret_cast<const bool *>(
+                                              windowReversalBool.data()),
+                                          windowReversalBool.size())),
+        convolutionDimensionNumbersAttr, featureGroupCount, batchGroupCount,
+        mlir::ArrayAttr::get(context, llvm::ArrayRef(precisionVec)));
+    break;
+  }
   case tensat::Ops::DotGeneralOp: {
     std::vector<int64_t> lhs_batch_dim = other_vecs[0];
     std::vector<int64_t> rhs_batch_dim = other_vecs[1];
@@ -651,7 +979,8 @@ Operation *createStableHloOp(OpBuilder &builder, tensat::Ops op,
 // so duplicated)
 uint64_t tensat::get_cost(tensat::Ops op, rust::Vec<tensat::Tensor> enode_args,
                           rust::Vec<tensat::Vector> other_vector_args,
-                          rust::Vec<int64_t> int_args) {
+                          rust::Vec<int64_t> int_args,
+                          rust::Vec<tensat::Matrix> matrix_args) {
   auto context = OperationTimer::getContext();
   OpBuilder builder(context);
 
@@ -671,9 +1000,23 @@ uint64_t tensat::get_cost(tensat::Ops op, rust::Vec<tensat::Tensor> enode_args,
   for (const auto &num : int_args)
     int_args_as_vec.push_back(num);
 
+  std::vector<std::vector<std::vector<int64_t>>> matrix_args_as_vec;
+
+  for (const auto &mat : matrix_args) {
+    std::vector<std::vector<int64_t>> current_matrix;
+
+    for (const auto &vec : mat.mat) {
+      std::vector<int64_t> converted_vec(vec.vec.begin(), vec.vec.end());
+      current_matrix.push_back(converted_vec);
+    }
+
+    matrix_args_as_vec.push_back(current_matrix);
+  }
+
   // Create the MLIR operation
-  Operation *mlirOp = createStableHloOp(builder, op, operands, other_vecs,
-                                        int_args_as_vec, context);
+  Operation *mlirOp =
+      createStableHloOp(builder, op, operands, other_vecs, int_args_as_vec,
+                        matrix_args_as_vec, context);
 
   int repeats = 0;
   switch (getPlatform()) {
@@ -743,8 +1086,9 @@ std::vector<int32_t> castArrayRefToInt32(llvm::ArrayRef<int64_t> shape) {
   return dims;
 }
 
-rust::Vec<int64_t> castArrayRefToRustVec(llvm::ArrayRef<int64_t> vec) {
-  rust::Vec<int64_t> res;
+template <typename T>
+rust::Vec<T> castArrayRefToRustVec(llvm::ArrayRef<T> vec) {
+  rust::Vec<T> res;
   res.reserve(vec.size());
   for (const auto &elem : vec) {
     res.push_back(elem);
@@ -752,11 +1096,55 @@ rust::Vec<int64_t> castArrayRefToRustVec(llvm::ArrayRef<int64_t> vec) {
   return res;
 }
 
+rust::Vec<tensat::Vector>
+castDenseIntElementsAttrToRustMatrix(mlir::DenseIntElementsAttr attr) {
+  rust::Vec<tensat::Vector> res;
+
+  // Get the shape of the tensor
+  auto shape = attr.getType().getShape();
+
+  if (shape.size() == 2) {
+    // 2D matrix case
+    auto numRows = shape[0];
+    auto numCols = shape[1];
+
+    // Iterator for the elements
+    auto it = attr.value_begin<llvm::APInt>();
+
+    for (int i = 0; i < numRows; ++i) {
+      tensat::Vector rowVector;
+      rowVector.vec.reserve(numCols);
+      for (int j = 0; j < numCols; ++j) {
+        // Extract the integer value from APInt and push it into the row vector
+        rowVector.vec.push_back((*it).getSExtValue());
+        ++it;
+      }
+      res.push_back(
+          std::move(rowVector)); // Push the row wrapped in Vector struct
+    }
+  } else if (shape.size() == 1) {
+    // 1D vector case
+    tensat::Vector rowVector;
+    for (auto it = attr.value_begin<llvm::APInt>();
+         it != attr.value_end<llvm::APInt>(); ++it) {
+      rowVector.vec.push_back(
+          (*it).getSExtValue()); // Push elements into the row vector
+    }
+    res.push_back(
+        std::move(rowVector)); // Push the 1D vector wrapped in Vector struct
+  } else {
+    llvm::errs() << "Unhandled tensor rank for DenseElementsAttr.";
+  }
+
+  return res;
+}
+
 // SHAPE INFERENCE
 rust::Vec<tensat::Tensor>
 tensat::get_shape(Ops op, rust::Vec<tensat::Tensor> enode_args,
                   rust::Vec<tensat::Vector> other_vector_args,
-                  rust::Vec<int64_t> int_args) {
+                  rust::Vec<int64_t> int_args,
+                  rust::Vec<tensat::Matrix> matrix_args) {
   auto context = OperationTimer::getContext();
   OpBuilder builder(context);
 
@@ -776,9 +1164,20 @@ tensat::get_shape(Ops op, rust::Vec<tensat::Tensor> enode_args,
   for (const auto &num : int_args)
     int_args_as_vec.push_back(num);
 
+  std::vector<std::vector<std::vector<int64_t>>> matrix_args_as_vec;
+  for (const auto &mat : matrix_args) {
+    std::vector<std::vector<int64_t>> current_matrix;
+    for (const auto &vec : mat.mat) {
+      std::vector<int64_t> converted_vec(vec.vec.begin(), vec.vec.end());
+      current_matrix.push_back(converted_vec);
+    }
+    matrix_args_as_vec.push_back(current_matrix);
+  }
+
   // Create the MLIR operation
-  Operation *mlirOp = createStableHloOp(builder, op, operands, other_vecs,
-                                        int_args_as_vec, context);
+  Operation *mlirOp =
+      createStableHloOp(builder, op, operands, other_vecs, int_args_as_vec,
+                        matrix_args_as_vec, context);
   if (mlirOp) {
     rust::Vec<tensat::Tensor> tensors;
     for (auto res : mlirOp->getResults()) {
@@ -1005,32 +1404,98 @@ class EqualitySaturationPass
         }
       }
 
-      if (auto output_tensor =
-              dot_general.getResult().getType().cast<TensorType>()) {
-        auto shape = castArrayRefToInt32(output_tensor.getShape());
-        auto output_shape_slice =
-            rust::Slice<const int>{shape.data(), shape.size()};
-
-        tensorInfo = graph
-                         ->new_dot_general_op(
-                             *handleOperandPartial(dot_general.getLhs()),
-                             *handleOperandPartial(dot_general.getRhs()),
-                             castArrayRefToRustVec(
-                                 dot_dim_attrs.getLhsBatchingDimensions()),
-                             castArrayRefToRustVec(
-                                 dot_dim_attrs.getRhsBatchingDimensions()),
-                             castArrayRefToRustVec(
-                                 dot_dim_attrs.getLhsContractingDimensions()),
-                             castArrayRefToRustVec(
-                                 dot_dim_attrs.getRhsContractingDimensions()),
-                             precision_configs,
-                             mlirValueToTensatTensor(dot_general.getResult()))
-                         .into_raw();
+      tensorInfo = graph
+                       ->new_dot_general_op(
+                           *handleOperandPartial(dot_general.getLhs()),
+                           *handleOperandPartial(dot_general.getRhs()),
+                           castArrayRefToRustVec(
+                               dot_dim_attrs.getLhsBatchingDimensions()),
+                           castArrayRefToRustVec(
+                               dot_dim_attrs.getRhsBatchingDimensions()),
+                           castArrayRefToRustVec(
+                               dot_dim_attrs.getLhsContractingDimensions()),
+                           castArrayRefToRustVec(
+                               dot_dim_attrs.getRhsContractingDimensions()),
+                           precision_configs,
+                           mlirValueToTensatTensor(dot_general.getResult()))
+                       .into_raw();
+    } else if (isa<stablehlo::ConvolutionOp>(op)) {
+      auto convolution = cast<stablehlo::ConvolutionOp>(op);
+      auto result_type =
+          convolution.getResult().getType().cast<mlir::ShapedType>();
+      if (result_type.hasRank()) {
+        auto shape = result_type.getShape();
+        llvm::errs() << "Result shape: [";
+        for (auto dim : shape) {
+          llvm::errs() << dim << " ";
+        }
+        llvm::errs() << "]\n";
       } else {
-        std::cout << "EqualitySaturationPass: result of "
-                     "stablehlo::DotGeneralOp has non-tensor type"
-                  << std::endl;
+        llvm::errs() << "Result is unranked.\n";
+      }
+      auto dimNumbers = convolution.getDimensionNumbers();
+      mlir::ArrayAttr precision =
+          convolution.getPrecisionConfig().value_or(mlir::ArrayAttr());
+      rust::Vec<int64_t> precision_configs;
+      for (int i = 0; i < precision.size(); i++) {
+        auto precisionAttr =
+            precision[i].dyn_cast<mlir::stablehlo::PrecisionAttr>();
+        if (!precisionAttr)
+          continue; // Skip if it's not a PrecisionAttr, although such
+                    // attributes should not exist here
+        mlir::stablehlo::Precision val = precisionAttr.getValue();
+        switch (val) {
+        case mlir::stablehlo::Precision::DEFAULT:
+          precision_configs.push_back(0);
+          break;
+        case mlir::stablehlo::Precision::HIGH:
+          precision_configs.push_back(1);
+          break;
+        case mlir::stablehlo::Precision::HIGHEST:
+          precision_configs.push_back(2);
+          break;
+        }
       }
+      auto windowStridesOpt = convolution.getWindowStrides();
+      assert(windowStridesOpt.has_value() &&
+             "Expected window strides to be present.");
+      auto paddingOpt = convolution.getPadding();
+      assert(paddingOpt.has_value() && "Expected padding to be present.");
+      auto lhsDilationOpt = convolution.getLhsDilation();
+      assert(lhsDilationOpt.has_value() &&
+             "Expected LHS dilation to be present.");
+      auto rhsDilationOpt = convolution.getRhsDilation();
+      assert(rhsDilationOpt.has_value() &&
+             "Expected RHS dilation to be present.");
+      auto windowReversalOpt = convolution.getWindowReversal();
+      assert(windowReversalOpt.has_value() &&
+             "Expected Window reversal to be present.");
+
+      tensorInfo =
+          graph
+              ->new_convolution_op(
+                  *handleOperandPartial(convolution.getLhs()),
+                  *handleOperandPartial(convolution.getRhs()),
+                  castArrayRefToRustVec(windowStridesOpt.value()),
+                  castDenseIntElementsAttrToRustMatrix(paddingOpt.value()),
+                  castArrayRefToRustVec(lhsDilationOpt.value()),
+                  castArrayRefToRustVec(rhsDilationOpt.value()),
+                  castArrayRefToRustVec(windowReversalOpt.value()),
+                  dimNumbers.getInputBatchDimension(),
+                  dimNumbers.getInputFeatureDimension(),
+                  castArrayRefToRustVec(dimNumbers.getInputSpatialDimensions()),
+                  dimNumbers.getKernelInputFeatureDimension(),
+                  dimNumbers.getKernelOutputFeatureDimension(),
+                  castArrayRefToRustVec(
+                      dimNumbers.getKernelSpatialDimensions()),
+                  dimNumbers.getOutputBatchDimension(),
+                  dimNumbers.getOutputFeatureDimension(),
+                  castArrayRefToRustVec(
+                      dimNumbers.getOutputSpatialDimensions()),
+                  convolution.getFeatureGroupCount(),
+                  convolution.getBatchGroupCount(), precision_configs,
+                  mlirValueToTensatTensor(convolution.getResult()))
+              .into_raw();
     } else if (isa<stablehlo::ConcatenateOp>(op)) {
       auto concat = cast<stablehlo::ConcatenateOp>(op);
       auto output_tensor = concat->getResult(0).getType().cast<TensorType>();
@@ -1088,8 +1553,8 @@ class EqualitySaturationPass
         outputs.push_back(mlirValueToTensatTensor(result));
 
       std::vector<tensat::TensorInfo *> processedOperands;
-      // We shouldn't clone operands, as those values will be invalidated after
-      // a block.clear(), and we access these during reconstruction.
+      // We shouldn't clone operands, as those values will be invalidated
+      // after a block.clear(), and we access these during reconstruction.
       auto copy = op->clone(Operation::CloneOptions(
           /* cloneRegions = */ true, /* cloneOperands = */ false));
       blackboxIDToTensorInfo->push_back(copy);
@@ -1180,25 +1645,43 @@ class EqualitySaturationPass
    */
   std::vector<int64_t> parseNumVec(rust::vec<tensat::Node> &nodes,
                                    tensat::Node &seq) {
-    assert(seq.name == "Vec");
+    assert(seq.op == tensat::Ops::Vec);
     std::vector<int64_t> result;
 
     for (auto i : seq.operands) {
       assert(i < nodes.size());
-      assert(nodes[i].name == "Num");
+      assert(nodes[i].op == tensat::Ops::Num);
       result.push_back(parseNumNode(nodes, nodes[i]));
     }
 
     return result;
   }
 
+  /**
+   * Parse the Vec nodes with Vecs (e.g Vec(Vec(128, 128), Vec(128, 128)))
+   * emitted by tensat node construction.
+   */
+  mlir::DenseIntElementsAttr
+  parseNumMatrixToDenseAttr(rust::vec<tensat::Node> &nodes, tensat::Node &seq,
+                            mlir::Builder &builder) {
+    assert(seq.op == tensat::Ops::Vec);
+
+    std::vector<std::vector<int64_t>> matrix;
+    for (auto i : seq.operands) {
+      assert(i < nodes.size());
+      assert(nodes[i].op == tensat::Ops::Vec);
+      matrix.push_back(parseNumVec(nodes, nodes[i]));
+    }
+    return matrixToDenseAttr(matrix, builder);
+  }
+
   /**
    * Parse the Num nodes emitted by tensat node construction.
    * Our protocol is to encode integer values as operand indices.
    * TODO: improve this!
    */
   int64_t parseNumNode(rust::vec<tensat::Node> &nodes, tensat::Node &seq) {
-    assert(seq.name == "Num");
+    assert(seq.op == tensat::Ops::Num);
     return seq.operands[0];
   }
 
@@ -1207,7 +1690,7 @@ class EqualitySaturationPass
    * AddOp(...))) emitted by tensat node construction.
    */
   std::vector<Value> parseOpVec(std::vector<Value> &opVals, tensat::Node &seq) {
-    assert(seq.name == "Vec");
+    assert(seq.op == tensat::Ops::Vec);
     std::vector<Value> result;
 
     for (auto i : seq.operands) {
@@ -1238,10 +1721,10 @@ class EqualitySaturationPass
     auto &region = funcOp.getRegion();
     auto &block = funcOp.getRegion().front();
 
-    // We don't clear the block here straight away, because this will invalidate
-    // the old captured values in blackbox, and so the substitution won't work.
-    // We instead put everything in a vector, and only clear the block and
-    // insert everything at the very end.
+    // We don't clear the block here straight away, because this will
+    // invalidate the old captured values in blackbox, and so the substitution
+    // won't work. We instead put everything in a vector, and only clear the
+    // block and insert everything at the very end.
     std::vector<Operation *> opsToAdd;
 
     auto location = builder.getUnknownLoc();
@@ -1321,6 +1804,120 @@ class EqualitySaturationPass
         }
         break;
       }
+      case Ops::InferReshape: {
+        auto input = opVals[node.operands[0]];
+        llvm::SmallVector<int64_t> shape;
+        for (int i = 1; i < node.operands.size(); i++)
+          shape.push_back(node.operands[i]);
+        if (shape == getShape(input)) {
+          opVals.push_back(input);
+          continue;
+        } else {
+          auto newType = deriveOutputType(input, shape);
+          newOp =
+              builder.create<stablehlo::ReshapeOp>(location, newType, input);
+          break;
+        }
+      }
+      case Ops::ConvolutionOp: {
+        auto lhs = opVals[node.operands[0]];
+        auto rhs = opVals[node.operands[1]];
+
+        auto windowStrides = parseNumVec(nodes, nodes[node.operands[2]]);
+        auto padding =
+            parseNumMatrixToDenseAttr(nodes, nodes[node.operands[3]], builder);
+        auto lhsDilation = parseNumVec(nodes, nodes[node.operands[4]]);
+        auto rhsDilation = parseNumVec(nodes, nodes[node.operands[5]]);
+        auto windowReversal = parseNumVec(nodes, nodes[node.operands[6]]);
+        std::vector<uint8_t> windowReversalBool;
+        windowReversalBool.reserve(windowReversal.size());
+        for (int64_t val : windowReversal) {
+          windowReversalBool.push_back(static_cast<uint8_t>(val != 0));
+        }
+        // dimension numbers
+        auto inputBatchDimension = parseNumNode(nodes, nodes[node.operands[7]]);
+        auto inputFeatureDimension =
+            parseNumNode(nodes, nodes[node.operands[8]]);
+        auto inputSpatialDimensions =
+            parseNumVec(nodes, nodes[node.operands[9]]);
+        auto kernelInputFeatureDimension =
+            parseNumNode(nodes, nodes[node.operands[10]]);
+        auto kernelOutputFeatureDimension =
+            parseNumNode(nodes, nodes[node.operands[11]]);
+        auto kernelSpatialDimensions =
+            parseNumVec(nodes, nodes[node.operands[12]]);
+        auto outputBatchDimension =
+            parseNumNode(nodes, nodes[node.operands[13]]);
+        auto outputFeatureDimension =
+            parseNumNode(nodes, nodes[node.operands[14]]);
+        auto outputSpatialDimensions =
+            parseNumVec(nodes, nodes[node.operands[15]]);
+        auto convolutionDimensionNumbersAttr =
+            stablehlo::ConvDimensionNumbersAttr::get(
+                context, inputBatchDimension, inputFeatureDimension,
+                inputSpatialDimensions, kernelInputFeatureDimension,
+                kernelOutputFeatureDimension, kernelSpatialDimensions,
+                outputBatchDimension, outputFeatureDimension,
+                outputSpatialDimensions);
+
+        ConvDimensionNumbers convDims;
+        convDims.lhs_spec = {inputBatchDimension, inputFeatureDimension};
+        convDims.lhs_spec.insert(convDims.lhs_spec.end(),
+                                 inputSpatialDimensions.begin(),
+                                 inputSpatialDimensions.end());
+
+        convDims.rhs_spec = {kernelOutputFeatureDimension,
+                             kernelInputFeatureDimension};
+        convDims.rhs_spec.insert(convDims.rhs_spec.end(),
+                                 kernelSpatialDimensions.begin(),
+                                 kernelSpatialDimensions.end());
+
+        convDims.out_spec = {outputBatchDimension, outputFeatureDimension};
+        convDims.out_spec.insert(convDims.out_spec.end(),
+                                 outputSpatialDimensions.begin(),
+                                 outputSpatialDimensions.end());
+        auto featureGroupCount = parseNumNode(nodes, nodes[node.operands[16]]);
+        auto batchGroupCount = parseNumNode(nodes, nodes[node.operands[17]]);
+        auto precisionConfig = parseNumVec(nodes, nodes[node.operands[18]]);
+
+        std::vector<Attribute> precisionVec;
+        for (auto &precision : precisionConfig) {
+          switch (precision) {
+          case 0:
+            precisionVec.push_back(stablehlo::PrecisionAttr::get(
+                context, stablehlo::Precision::DEFAULT));
+            break;
+          case 1:
+            precisionVec.push_back(stablehlo::PrecisionAttr::get(
+                context, stablehlo::Precision::HIGH));
+            break;
+          case 2:
+            precisionVec.push_back(stablehlo::PrecisionAttr::get(
+                context, stablehlo::Precision::HIGHEST));
+            break;
+          }
+        }
+
+        auto shape = convolutionShapeComputation(
+            getShape(lhs), getShape(rhs), windowStrides, padding, lhsDilation,
+            rhsDilation, convDims, featureGroupCount, batchGroupCount);
+
+        auto newType = deriveOutputType(lhs, shape);
+        newOp = builder.create<stablehlo::ConvolutionOp>(
+            location, newType, lhs, rhs,
+            mlir::DenseI64ArrayAttr::get(context,
+                                         llvm::ArrayRef(windowStrides)),
+            padding,
+            mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(lhsDilation)),
+            mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(rhsDilation)),
+            mlir::DenseBoolArrayAttr::get(
+                context, llvm::ArrayRef<bool>(reinterpret_cast<const bool *>(
+                                                  windowReversalBool.data()),
+                                              windowReversalBool.size())),
+            convolutionDimensionNumbersAttr, featureGroupCount, batchGroupCount,
+            mlir::ArrayAttr::get(context, llvm::ArrayRef(precisionVec)));
+        break;
+      }
       case Ops::DotGeneralOp: {
         auto lhs = opVals[node.operands[0]];
         auto rhs = opVals[node.operands[1]];
@@ -1437,8 +2034,6 @@ class EqualitySaturationPass
           subst.map(oldCapturedValues[i], capturedValues[i]);
         }
         newOp = newOp->clone(subst);
-
-        assert(numOperands == newOp->getNumOperands());
         newOp->insertOperands(0, operands);
         break;
       }
@@ -1507,8 +2102,8 @@ class EqualitySaturationPass
     SmallVector<Operation *> currentOps;
     SegmentationPoint segment;
 
-    // We need to keep track of anything that was an output value, so that if we
-    // see it in later segments then we know to add that as an input.
+    // We need to keep track of anything that was an output value, so that if
+    // we see it in later segments then we know to add that as an input.
     DenseSet<Value> outputsBeforeCurrentSegment;
 
     int nonBlackboxedInCurrentSegment = 0;
@@ -1531,8 +2126,8 @@ class EqualitySaturationPass
 
       bool blackboxed = isBlackboxed(&op);
 
-      // TODO: This ensures the last node in segment is blackboxed, but ideally
-      // we actually want to reduce the number of segment outputs.
+      // TODO: This ensures the last node in segment is blackboxed, but
+      // ideally we actually want to reduce the number of segment outputs.
       if ((blackboxed && nonBlackboxedInCurrentSegment >= segmentThreshold) ||
           it == (--entryBlock.end())) {
         // Track outputs of this segment
diff --git a/src/enzyme_ad/jax/Passes/EqualitySaturation.h b/src/enzyme_ad/jax/Passes/EqualitySaturation.h
index 96dc5450..1d52c9c0 100644
--- a/src/enzyme_ad/jax/Passes/EqualitySaturation.h
+++ b/src/enzyme_ad/jax/Passes/EqualitySaturation.h
@@ -9,6 +9,7 @@ namespace tensat {
 enum class Type : uint8_t;
 enum class Ops : uint8_t;
 struct Vector;
+struct Matrix;
 struct Tensor;
 
 /**
@@ -17,12 +18,14 @@ struct Tensor;
 
 uint64_t get_cost(Ops op, rust::Vec<tensat::Tensor> operands,
                   rust::Vec<tensat::Vector> other_vector_args,
-                  rust::Vec<int64_t> int_args);
+                  rust::Vec<int64_t> int_args,
+                  rust::Vec<tensat::Matrix> matrix_args);
 
 mlir::Type newTensorType(mlir::OpBuilder &builder, Tensor tensor); 
 mlir::Type tensatTypeToMlirType(mlir::OpBuilder &builder, Type type);
 
 rust::Vec<Tensor> get_shape(Ops op, rust::Vec<tensat::Tensor> operands,
                            rust::Vec<tensat::Vector> other_vector_args,
-                           rust::Vec<int64_t> int_args);
+                           rust::Vec<int64_t> int_args,
+                  rust::Vec<tensat::Matrix> matrix_args);
 } // namespace tensat
diff --git a/src/enzyme_ad/jax/TransformOps/TransformOps.td b/src/enzyme_ad/jax/TransformOps/TransformOps.td
index 79112780..0201fa41 100644
--- a/src/enzyme_ad/jax/TransformOps/TransformOps.td
+++ b/src/enzyme_ad/jax/TransformOps/TransformOps.td
@@ -125,10 +125,6 @@ def ApplyConvertSimplifyPatterns : EnzymeHLOPatternOp<
     "convert_simplify"> {
   let patterns = ["ConvertSimplify"];
 }
-def ApplyReshapeSimplifyPatterns : EnzymeHLOPatternOp<
-    "reshape_simplify"> {
-  let patterns = ["ReshapeSimplify"];
-}
 def ApplyDynamicSliceToStaticPatterns : EnzymeHLOPatternOp<
     "dynamic_slice_to_static"> {
   let patterns = ["DynamicSliceToStatic"];
diff --git a/src/enzyme_ad/jax/deps/tensat/Cargo.Bazel.lock b/src/enzyme_ad/jax/deps/tensat/Cargo.Bazel.lock
index dd319354..063b4d12 100644
--- a/src/enzyme_ad/jax/deps/tensat/Cargo.Bazel.lock
+++ b/src/enzyme_ad/jax/deps/tensat/Cargo.Bazel.lock
@@ -207,7 +207,7 @@ dependencies = [
 [[package]]
 name = "egg"
 version = "0.6.1-dev"
-source = "git+https://github.com/yycdavid/egg?rev=12cc1ee7731d37fe91901c81f59678fa1d08a2bb#12cc1ee7731d37fe91901c81f59678fa1d08a2bb"
+source = "git+https://github.com/aryavohra/egg?rev=b30d14cff61bff97336323f6eb0978cc7769140d#b30d14cff61bff97336323f6eb0978cc7769140d"
 dependencies = [
  "indexmap",
  "instant",
diff --git a/src/enzyme_ad/jax/deps/tensat/Cargo.lock b/src/enzyme_ad/jax/deps/tensat/Cargo.lock
index 9e3810fb..78a76be3 100644
--- a/src/enzyme_ad/jax/deps/tensat/Cargo.lock
+++ b/src/enzyme_ad/jax/deps/tensat/Cargo.lock
@@ -207,7 +207,7 @@ dependencies = [
 [[package]]
 name = "egg"
 version = "0.6.1-dev"
-source = "git+https://github.com/yycdavid/egg?rev=12cc1ee7731d37fe91901c81f59678fa1d08a2bb#12cc1ee7731d37fe91901c81f59678fa1d08a2bb"
+source = "git+https://github.com/aryavohra/egg?rev=b30d14cff61bff97336323f6eb0978cc7769140d#b30d14cff61bff97336323f6eb0978cc7769140d"
 dependencies = [
  "indexmap",
  "instant",
diff --git a/src/enzyme_ad/jax/deps/tensat/Cargo.toml b/src/enzyme_ad/jax/deps/tensat/Cargo.toml
index 93242369..11d431f9 100644
--- a/src/enzyme_ad/jax/deps/tensat/Cargo.toml
+++ b/src/enzyme_ad/jax/deps/tensat/Cargo.toml
@@ -25,8 +25,8 @@ serde_json = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 
 [dependencies.egg]
-git = "https://github.com/yycdavid/egg"
-rev = "12cc1ee7731d37fe91901c81f59678fa1d08a2bb"
+git = "https://github.com/aryavohra/egg"
+rev = "b30d14cff61bff97336323f6eb0978cc7769140d"
 
 [package.metadata.cxx]
 library = "c++"
diff --git a/src/enzyme_ad/jax/deps/tensat/converted.txt b/src/enzyme_ad/jax/deps/tensat/converted.txt
index 36b9a906..c4116ada 100644
--- a/src/enzyme_ad/jax/deps/tensat/converted.txt
+++ b/src/enzyme_ad/jax/deps/tensat/converted.txt
@@ -26,3 +26,11 @@
 (ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?z ?w)) ?i)<=>(MulOp (ConcatenateOp (Vec ?x ?z) ?i) (ConcatenateOp (Vec ?y ?w) ?i))
 
 (ConcatenateOp (Vec (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?z ?w) 1)) 0)<=>(ConcatenateOp (Vec (ConcatenateOp (Vec ?x ?z) 0) (ConcatenateOp (Vec ?y ?w) 0)) 1)
+
+(ConvolutionOp (MulOp ?x ?w) ?y ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)<=>(ConvolutionOp ?x (MulOp ?y ?w) ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)
+
+(ConvolutionOp ?lhs (ConcatenateOp (Vec ?x ?y) ?i) ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)<=>(ConcatenateOp (Vec (ConvolutionOp ?lhs ?x ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig) (ConvolutionOp ?lhs ?y ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)) ?i)
+
+(ConvolutionOp (AddOp ?x ?y) ?rhs ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)<=>(AddOp (ConvolutionOp ?x ?rhs ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig) (ConvolutionOp ?y ?rhs ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig))
+
+(ConvolutionOp ?lhs (AddOp ?x ?y) ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)<=>(AddOp (ConvolutionOp ?lhs ?x ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig) (ConvolutionOp ?lhs ?y ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig))
diff --git a/src/enzyme_ad/jax/deps/tensat/converted_multi-original.txt b/src/enzyme_ad/jax/deps/tensat/converted_multi-original.txt
index 8f608b45..eac8a812 100644
--- a/src/enzyme_ad/jax/deps/tensat/converted_multi-original.txt
+++ b/src/enzyme_ad/jax/deps/tensat/converted_multi-original.txt
@@ -23,4 +23,4 @@
 (conv2d 1 1 0 0 ?input_7 ?input_10)=>(split_0 (split 1 (conv2d 1 1 0 0 ?input_7 (concat 0 4 ?input_10 ?input_11))))
 (conv2d 1 1 0 0 ?input_7 ?input_11)=>(split_1 (split 1 (conv2d 1 1 0 0 ?input_7 (concat 0 4 ?input_10 ?input_11))))
 (relu (conv2d 1 1 0 0 ?input_7 ?input_10))=>(split_0 (split 1 (conv2d 1 1 0 2 ?input_7 (concat 0 4 ?input_10 ?input_11))))
-(conv2d 1 1 0 2 ?input_7 ?input_11)=>(split_1 (split 1 (conv2d 1 1 0 2 ?input_7 (concat 0 4 ?input_10 ?input_11))))
\ No newline at end of file
+(conv2d 1 1 0 2 ?input_7 ?input_11)=>(split_1 (split 1 (conv2d 1 1 0 2 ?input_7 (concat 0 4 ?input_10 ?input_11))))
diff --git a/src/enzyme_ad/jax/deps/tensat/converted_multi.txt b/src/enzyme_ad/jax/deps/tensat/converted_multi.txt
index ce117f3a..db259805 100644
--- a/src/enzyme_ad/jax/deps/tensat/converted_multi.txt
+++ b/src/enzyme_ad/jax/deps/tensat/converted_multi.txt
@@ -1,4 +1,15 @@
 (DotGeneralOp ?x ?y (Vec) (Vec) (Vec 1) (Vec 0) ?p)=>(MatchRank (SSplit0 (DotGeneralOp ?x (ConcatenateOp (Vec (MatchRank ?y ?x) (MatchRank ?z ?x)) 1) (Vec) (Vec) (Vec 1) (Vec 0) ?p) 1 (MatchRank ?y ?x)) ?y)
 (DotGeneralOp ?x ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p)=>(MatchRank (SSplit1 (DotGeneralOp ?x (ConcatenateOp (Vec (MatchRank ?y ?x) (MatchRank ?z ?x)) 1) (Vec) (Vec) (Vec 1) (Vec 0) ?p) 1 (MatchRank ?z ?x)) ?z)
+
+(DotGeneralOp ?x ?y (Vec) (Vec) (Vec 2) (Vec 0) ?p)=>(InferReshape (SSplit0 (DotGeneralOp ?x (ConcatenateOp (Vec (MatchRank ?y ?x) (MatchRank ?z ?x)) 2) (Vec) (Vec) (Vec 2) (Vec 0) ?p) 2 (MatchRank ?y ?x)))
+(DotGeneralOp ?x ?z (Vec) (Vec) (Vec 2) (Vec 0) ?p)=>(InferReshape (SSplit1 (DotGeneralOp ?x (ConcatenateOp (Vec (MatchRank ?y ?x) (MatchRank ?z ?x)) 2) (Vec) (Vec) (Vec 2) (Vec 0) ?p) 2 (MatchRank ?z ?x)))
+
 (DotGeneralOp ?x ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p)=>(SSplit0 (DotGeneralOp (ConcatenateOp (Vec ?x ?y) 0) ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p) 0 ?x)
 (DotGeneralOp ?y ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p)=>(SSplit1 (DotGeneralOp (ConcatenateOp (Vec ?x ?y) 0) ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p) 0 ?y)
+
+(ConvolutionOp ?x ?rhs ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)=>(SSplit0 (ConvolutionOp (ConcatenateOp (Vec ?x ?y) 0) ?rhs ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig) 0 ?x)
+(ConvolutionOp ?y ?rhs ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)=>(SSplit1 (ConvolutionOp (ConcatenateOp (Vec ?x ?y) 0) ?rhs ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig) 0 ?y)
+
+(ConvolutionOp ?lhs ?x ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)=>(SSplit0 (ConvolutionOp ?lhs (ConcatenateOp (Vec ?x ?y) 0) ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig) 0 ?x)
+(ConvolutionOp ?lhs ?y ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig)=>(SSplit1 (ConvolutionOp ?lhs (ConcatenateOp (Vec ?x ?y) 0) ?windowstrides ?padding ?lhsdilation ?rhsdilation ?windowreversal ?inputbatchdimension ?inputfeaturedimension ?inputspatialdimensions ?kernelinputfeaturedimension ?kerneloutputfeaturedimension ?kernelspatialdimensions ?outputbatchdimension ?outputfeaturedimension ?outputspatialdimensions ?featuregroupcount ?batchgroupcount ?precisionconfig) 0 ?y)
+
diff --git a/src/enzyme_ad/jax/deps/tensat/src/ffi_utils.rs b/src/enzyme_ad/jax/deps/tensat/src/ffi_utils.rs
index 9ed2e6de..d3da0ac5 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/ffi_utils.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/ffi_utils.rs
@@ -1,21 +1,29 @@
 use crate::{
     input::ffi,
     model::*,
-    rewrites::{get_num_option, get_vec_of_nums_option, get_vec_option},
+    rewrites::{get_matrix_option, get_num_option, get_vec_of_nums_option, get_vec_option},
 };
 use egg::*;
 
 fn process_enode_args(
     egraph: &EGraph<Mdl, TensorAnalysis>,
     enode: &Mdl,
-) -> (Vec<ffi::Tensor>, Vec<ffi::Vector>, Vec<i64>) {
+) -> (
+    Vec<ffi::Tensor>,
+    Vec<ffi::Vector>,
+    Vec<i64>,
+    Vec<ffi::Matrix>,
+) {
     let mut args: Vec<ffi::Tensor> = vec![];
     let mut other_vecs: Vec<ffi::Vector> = vec![];
     let mut int_args: Vec<i64> = vec![];
+    let mut matrix_args: Vec<ffi::Matrix> = vec![];
 
     for child in enode.children().iter() {
         if let Some(other_vec) = get_vec_of_nums_option(egraph, &egraph[*child]) {
             other_vecs.push(other_vec)
+        } else if let Some(mat) = get_matrix_option(egraph, &egraph[*child]) {
+            matrix_args.push(mat)
         } else if let Some(vec) = get_vec_option(&egraph[*child]) {
             vec.iter()
                 .for_each(|&id| args.push(egraph[id].data.tensors[0].clone()))
@@ -27,7 +35,7 @@ fn process_enode_args(
         }
     }
 
-    (args, other_vecs, int_args)
+    (args, other_vecs, int_args, matrix_args)
 }
 
 pub fn create_stablehlo_op<F, R>(
@@ -36,10 +44,10 @@ pub fn create_stablehlo_op<F, R>(
     process_output: F,
 ) -> R
 where
-    F: Fn(ffi::Ops, Vec<ffi::Tensor>, Vec<ffi::Vector>, Vec<i64>) -> R,
+    F: Fn(ffi::Ops, Vec<ffi::Tensor>, Vec<ffi::Vector>, Vec<i64>, Vec<ffi::Matrix>) -> R,
 {
     let op = ffi::Ops::from_mdl(enode);
-    let (args, other_vecs, int_args) = process_enode_args(egraph, enode);
-    let res = process_output(op, args, other_vecs, int_args);
+    let (args, other_vecs, int_args, matrix_args) = process_enode_args(egraph, enode);
+    let res = process_output(op, args, other_vecs, int_args, matrix_args);
     res
 }
diff --git a/src/enzyme_ad/jax/deps/tensat/src/input.rs b/src/enzyme_ad/jax/deps/tensat/src/input.rs
index 422e9546..73aadcae 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/input.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/input.rs
@@ -34,6 +34,7 @@ pub mod ffi {
         SelectOp,
         ConcatenateOp,
         DotGeneralOp,
+        ConvolutionOp,
         PadOp,
         SliceOp,
         TransposeOp,
@@ -56,6 +57,7 @@ pub mod ffi {
         SSplit0,
         SSplit1,
         MatchRank,
+        InferReshape,
     }
 
     struct Node {
@@ -83,6 +85,12 @@ pub mod ffi {
         pub vec: Vec<i64>,
     }
 
+    // Similarly, we're creating a Matrix type for vecs of vecs (padding)
+    #[derive(Debug)]
+    struct Matrix {
+        pub mat: Vec<Vector>,
+    }
+
     // take floats from c++ and wrap them into f32s below
     extern "Rust" {
         type Mdl;
@@ -159,6 +167,29 @@ pub mod ffi {
             dimension: i64,
             output: Tensor,
         ) -> Box<TensorInfo>;
+        fn new_convolution_op(
+            self: &mut CppGraphConverter,
+            lhs: &TensorInfo,
+            rhs: &TensorInfo,
+            windowStrides: Vec<i64>,
+            padding: Vec<Vector>,
+            lhsDilation: Vec<i64>,
+            rhsDilation: Vec<i64>,
+            windowReversal: Vec<bool>,
+            inputBatchDimension: i64,
+            inputFeatureDimension: i64,
+            inputSpatialDimension: Vec<i64>,
+            kernelInputFeatureDimension: i64,
+            kernelOutputFeatureDimension: i64,
+            kernelSpatialDimension: Vec<i64>,
+            outputBatchDimension: i64,
+            outputFeatureDimension: i64,
+            outputSpatialDimension: Vec<i64>,
+            featureGroupCount: i64,
+            batchGroupCount: i64,
+            precision_config: Vec<i64>,
+            output: Tensor,
+        ) -> Box<TensorInfo>;
         fn new_dot_general_op(
             self: &mut CppGraphConverter,
             lhs: &TensorInfo,
@@ -274,7 +305,7 @@ pub mod ffi {
         fn new_blackbox_op(
             self: &mut CppGraphConverter,
             inpts: &[*mut TensorInfo],
-            captured: &[*mut TensorInfo],  // values that appear in a block that was declared outside
+            captured: &[*mut TensorInfo], // values that appear in a block that was declared outside
             cpp_num: i64,
             outputs: &Vec<Tensor>,
         ) -> Box<TensorInfo>;
@@ -293,6 +324,7 @@ pub mod ffi {
             operands: Vec<Tensor>,
             other_vector_args: Vec<Vector>,
             int_args: Vec<i64>,
+            matrix_args: Vec<Matrix>,
         ) -> u64;
     }
 
@@ -304,6 +336,7 @@ pub mod ffi {
             operands: Vec<Tensor>,
             other_vector_args: Vec<Vector>,
             int_args: Vec<i64>,
+            matrix_args: Vec<Matrix>,
         ) -> Vec<Tensor>;
     }
 }
@@ -356,6 +389,7 @@ impl ffi::Ops {
             Mdl::PadOp(_) => Ops::PadOp,
             Mdl::SliceOp(_) => Ops::SliceOp,
             Mdl::TransposeOp(_) => Ops::TransposeOp,
+            Mdl::ConvolutionOp(_) => Ops::ConvolutionOp,
             Mdl::MulOp(_) => Ops::MulOp,
             Mdl::AddOp(_) => Ops::AddOp,
             Mdl::DivOp(_) => Ops::DivOp,
@@ -374,6 +408,7 @@ impl ffi::Ops {
             Mdl::SSplit0(_) => Ops::SSplit0,
             Mdl::SSplit1(_) => Ops::SSplit1,
             Mdl::MatchRank(_) => Ops::MatchRank,
+            Mdl::InferReshape(_) => Ops::InferReshape,
         }
     }
 }
@@ -404,6 +439,14 @@ impl CppGraphConverter {
         }
     }
 
+    fn tensor_data(tensors: Vec<ffi::Tensor>) -> TensorData {
+        TensorData {
+            tensors,
+            name: None,
+            need_infer_shape: false,
+        }
+    }
+
     // Wrapper functions for C++ side
     pub fn new_input(&mut self, block_arg_number: i64, tensor: ffi::Tensor) -> Box<TensorInfo> {
         let name = format!("input_{}", block_arg_number)
@@ -417,10 +460,7 @@ impl CppGraphConverter {
         let new_node = Mdl::Input([name_id, block_arg_node_id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![tensor],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![tensor]),
         };
         Box::new(res)
     }
@@ -430,10 +470,7 @@ impl CppGraphConverter {
         let new_node = Mdl::Index([index_num_node, inpt.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![inpt.tensor_data.tensors[index as usize].clone()],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![inpt.tensor_data.tensors[index as usize].clone()]),
         };
         Box::new(res)
     }
@@ -457,10 +494,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -476,10 +510,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -495,10 +526,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -514,10 +542,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: outputs.clone(),
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(outputs.clone()),
         };
         Box::new(res)
     }
@@ -528,10 +553,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -574,10 +596,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -593,18 +612,12 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
 
-    fn new_tensorinfo_vec(
-        &mut self,
-        inputs: &[*mut TensorInfo]
-    ) -> Id {
+    fn new_tensorinfo_vec(&mut self, inputs: &[*mut TensorInfo]) -> Id {
         let tensor_infos: Vec<&TensorInfo> = inputs.iter().map(|&ptr| unsafe { &*ptr }).collect();
         let inputs_node = Mdl::Vec(tensor_infos.iter().map(|i| i.id).collect());
         self.rec_expr.add(inputs_node)
@@ -622,10 +635,77 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
+        };
+        Box::new(res)
+    }
+
+    pub fn new_convolution_op(
+        &mut self,
+        lhs: &TensorInfo,
+        rhs: &TensorInfo,
+        window_strides: Vec<i64>,
+        padding: Vec<ffi::Vector>,
+        lhs_dilation: Vec<i64>,
+        rhs_dilation: Vec<i64>,
+        window_reversal: Vec<bool>,
+        input_batch_dimension: i64,
+        input_feature_dimension: i64,
+        input_spatial_dimensions: Vec<i64>,
+        kernel_input_feature_dimension: i64,
+        kernel_output_feature_dimension: i64,
+        kernel_spatial_dimensions: Vec<i64>,
+        output_batch_dimension: i64,
+        output_feature_dimension: i64,
+        output_spatial_dimensions: Vec<i64>,
+        feature_group_count: i64,
+        batch_group_count: i64,
+        precision_config: Vec<i64>,
+        output: ffi::Tensor,
+    ) -> Box<TensorInfo> {
+        let window_strides_node_id = self.vec_node(window_strides);
+        let lhs_dilation_node_id = self.vec_node(lhs_dilation);
+        let rhs_dilation_node_id = self.vec_node(rhs_dilation);
+
+        // We could add a bool element type vec?
+        let window_reversal_node_id =
+            self.vec_node(window_reversal.iter().map(|x| *x as i64).collect());
+        let input_spatial_dimensions_node_id = self.vec_node(input_spatial_dimensions);
+        let kernel_spatial_dimensions_node_id = self.vec_node(kernel_spatial_dimensions);
+        let output_spatial_dimensions_node_id = self.vec_node(output_spatial_dimensions);
+        let precision_config_node_id = self.vec_node(precision_config);
+
+        let padding_node_ids: Vec<Id> = padding
+            .into_iter()
+            .map(|pad| self.vec_node(pad.vec))
+            .collect::<Vec<Id>>();
+        let padding_node_id = self.rec_expr.add(Mdl::Vec(padding_node_ids));
+
+        let new_node = Mdl::ConvolutionOp([
+            lhs.id,
+            rhs.id,
+            window_strides_node_id,
+            padding_node_id,
+            lhs_dilation_node_id,
+            rhs_dilation_node_id,
+            window_reversal_node_id,
+            self.add_or_get_val(input_batch_dimension),
+            self.add_or_get_val(input_feature_dimension),
+            input_spatial_dimensions_node_id,
+            self.add_or_get_val(kernel_input_feature_dimension),
+            self.add_or_get_val(kernel_output_feature_dimension),
+            kernel_spatial_dimensions_node_id,
+            self.add_or_get_val(output_batch_dimension),
+            self.add_or_get_val(output_feature_dimension),
+            output_spatial_dimensions_node_id,
+            self.add_or_get_val(feature_group_count),
+            self.add_or_get_val(batch_group_count),
+            precision_config_node_id,
+        ]);
+
+        let res = TensorInfo {
+            id: self.rec_expr.add(new_node),
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -659,10 +739,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -690,10 +767,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -713,10 +787,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -732,10 +803,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -749,10 +817,7 @@ impl CppGraphConverter {
         let new_node = Mdl::MulOp([lhs.id, rhs.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -766,10 +831,7 @@ impl CppGraphConverter {
         let new_node = Mdl::AddOp([lhs.id, rhs.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -783,10 +845,7 @@ impl CppGraphConverter {
         let new_node = Mdl::DivOp([lhs.id, rhs.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -800,10 +859,7 @@ impl CppGraphConverter {
         let new_node = Mdl::SubtractOp([lhs.id, rhs.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -817,10 +873,7 @@ impl CppGraphConverter {
         let new_node = Mdl::MinOp([lhs.id, rhs.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -834,10 +887,7 @@ impl CppGraphConverter {
         let new_node = Mdl::MaxOp([lhs.id, rhs.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -846,10 +896,7 @@ impl CppGraphConverter {
         let new_node = Mdl::NegOp([inpt.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -858,10 +905,7 @@ impl CppGraphConverter {
         let new_node = Mdl::TanhOp([inpt.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -870,10 +914,7 @@ impl CppGraphConverter {
         let new_node = Mdl::ExpOp([inpt.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -884,10 +925,7 @@ impl CppGraphConverter {
         let new_node = Mdl::IotaOp([iota_dim_id, shape_id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -902,10 +940,7 @@ impl CppGraphConverter {
         let new_node = Mdl::DynamicUpdateSliceOp([operand.id, update.id, start_indices.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -922,10 +957,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![output],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![output]),
         };
         Box::new(res)
     }
@@ -947,10 +979,7 @@ impl CppGraphConverter {
         ]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: outputs.clone(),
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(outputs.clone()),
         };
         Box::new(res)
     }
@@ -969,10 +998,7 @@ impl CppGraphConverter {
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: outputs.clone(),
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(outputs.clone()),
         };
         self.blackbox_cpp_num_to_tensorinfo
             .insert(cpp_num, res.clone());
@@ -985,10 +1011,7 @@ impl CppGraphConverter {
         // Returns do not produce values!
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: TensorData {
-                tensors: vec![],
-                name: None,
-            },
+            tensor_data: CppGraphConverter::tensor_data(vec![]),
         };
         Box::new(res)
     }
@@ -1001,7 +1024,12 @@ impl CppGraphConverter {
         println!("{}", self.rec_expr.pretty(width as usize))
     }
 
-    fn convert_to_node(&self, rec_expr: RecExpr<Mdl>) -> Vec<ffi::Node> {
+    fn convert_to_node(
+        &self,
+        egraph: &EGraph<Mdl, TensorAnalysis>,
+        to_egraph: &HashMap<Id, Id>,
+        rec_expr: RecExpr<Mdl>
+    ) -> Vec<ffi::Node> {
         let mut res: Vec<ffi::Node> = Vec::new();
 
         let index = |id: Id| (usize::from(id) as i32); // TODO: this is probably wrong
@@ -1014,11 +1042,11 @@ impl CppGraphConverter {
 
         let rec_expr_ref = rec_expr.as_ref();
 
-        for mdl in rec_expr_ref.iter() {
+        for (i, mdl) in rec_expr_ref.iter().enumerate() {
             let op = ffi::Ops::from_mdl(mdl);
 
             let new_node = |operands: &[Id]| ffi::Node {
-                op: op,
+                op,
                 label: "".to_string(),
                 operands: convert(operands),
             };
@@ -1030,10 +1058,23 @@ impl CppGraphConverter {
                     operands: vec![],
                 },
                 Mdl::Num(num) => ffi::Node {
-                    op: op,
+                    op,
                     label: "".to_string(),
                     operands: vec![*num as i32],
                 },
+                Mdl::InferReshape([input]) => {
+                    let input_index = index(*input);
+                    let id = to_egraph[&Id::from(i)];
+                    let mut operands: Vec<i32> =
+                        (&egraph[id]).data.tensors[0].shape
+                            .iter().map(|x| *x as i32).collect();
+                    operands.insert(0, input_index);
+                    ffi::Node {
+                        op,
+                        label: "".to_string(),
+                        operands,
+                    }
+                }
                 // TODO: More clever pattern matching
                 Mdl::Vec(ops) => new_node(ops),
                 Mdl::Input(ops) => new_node(ops),
@@ -1043,6 +1084,7 @@ impl CppGraphConverter {
                 Mdl::DotGeneralOp(ops) => new_node(ops),
                 Mdl::SliceOp(ops) => new_node(ops),
                 Mdl::TransposeOp(ops) => new_node(ops),
+                Mdl::ConvolutionOp(ops) => new_node(ops),
                 Mdl::MulOp(ops) => new_node(ops),
                 Mdl::AddOp(ops) => new_node(ops),
                 Mdl::DivOp(ops) => new_node(ops),
@@ -1059,7 +1101,7 @@ impl CppGraphConverter {
                 Mdl::SSplit0(ops) => new_node(ops),
                 Mdl::SSplit1(ops) => new_node(ops),
                 Mdl::MatchRank(ops) => new_node(ops),
-                _ => unimplemented!()
+                _ => unimplemented!(),
             };
 
             res.push(node);
@@ -1087,10 +1129,9 @@ impl CppGraphConverter {
         let learned_rules =
             read_to_string(rule_file).expect("Something went wrong reading the rule file");
         let time_limit_sec = Duration::new(n_sec, 0);
-        let pre_defined_rules = PRE_DEFINED_RULES.iter().map(|&x| x);
-        let split_rules: Vec<&str> = learned_rules.split("\n")
+        let split_rules: Vec<&str> = learned_rules
+            .split("\n")
             .filter(|x| !x.is_empty())
-            .chain(pre_defined_rules)
             .collect();
         let do_filter_after = no_cycle && filter_after;
         let analysis = TensorAnalysis::new(&self.blackbox_cpp_num_to_tensorinfo);
@@ -1179,11 +1220,11 @@ impl CppGraphConverter {
 
         let (egraph, root) = (runner.egraph, runner.roots[0]);
         let cost_model: CostModel = CostModel::new();
-        let (best, ext_secs) = extract_by_ilp(&egraph, root, &cost_model);
+        let (best, ext_secs, to_egraph) = extract_by_ilp(&egraph, root, &cost_model);
         // let (best, ext_secs) = extract_by_greedy(&egraph, root, &cost_model);
 
         // println!("{}", best);
-        self.convert_to_node(best)
+        self.convert_to_node(&egraph, &to_egraph, best)
     }
 }
 
@@ -1210,7 +1251,7 @@ fn extract_by_ilp(
     egraph: &EGraph<Mdl, TensorAnalysis>,
     root: Id,
     cost_model: &CostModel,
-) -> (RecExpr<Mdl>, f32) {
+) -> (RecExpr<Mdl>, f32, HashMap<Id, Id>) {
     // Prepare data for ILP formulation, save to json
     let (m_id_map, e_m, h_i, cost_i, fus_i, g_i, root_m, i_to_nodes, blacklist_i) =
         prep_ilp_data(egraph, root, cost_model);
@@ -1234,7 +1275,10 @@ fn extract_by_ilp(
     let class_constraint = true;
     let no_order = true;
     let initialise_with_greedy = false;
-    let fusion_costs: bool = std::env::var("FUSION_COSTS").unwrap_or(String::from("true")).parse().unwrap();
+    let fusion_costs: bool = std::env::var("FUSION_COSTS")
+        .unwrap_or(String::from("false"))
+        .parse()
+        .unwrap();
     let mut arg_vec = vec!["src/enzyme_ad/jax/deps/tensat/extractor/extract.py"];
     if order_var_int {
         arg_vec.push("--order_var_int");
@@ -1281,7 +1325,8 @@ fn extract_by_ilp(
     arg_vec.push(time_lim);
     arg_vec.push("--num_thread");
     arg_vec.push(num_thread);
-    let child = Command::new("python")
+
+    let child = Command::new("python3")
         .args(&arg_vec)
         .spawn()
         .expect("failed to execute child");
@@ -1311,8 +1356,9 @@ fn extract_by_ilp(
 
         let mut expr = RecExpr::default();
         let mut added_memo: HashMap<Id, Id> = Default::default();
-        let _ = construct_best_rec(&node_picked, root, &mut added_memo, egraph, &mut expr);
-        (expr, solved_data.time)
+        let mut to_egraph: HashMap<Id, Id> = Default::default();
+        let _ = construct_best_rec(&node_picked, root, &mut added_memo, &mut to_egraph, egraph, &mut expr);
+        (expr, solved_data.time, to_egraph)
     } else {
         panic!("Python script failed");
     }
diff --git a/src/enzyme_ad/jax/deps/tensat/src/model.rs b/src/enzyme_ad/jax/deps/tensat/src/model.rs
index 316371f3..0cae0690 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/model.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/model.rs
@@ -26,6 +26,7 @@ define_language! {
       "GatherOp"           = GatherOp([Id; 10]),
       "SelectOp"           = SelectOp([Id; 3]), // pred, on_true, on_false
       "ConcatenateOp"      = ConcatenateOp([Id; 2]), // inputs, dimension
+      "ConvolutionOp"      = ConvolutionOp([Id; 19]), // LOTS of inputs
       "DotGeneralOp"       = DotGeneralOp([Id; 7]), // lhs, rhs, ..., shape
       "PadOp"              = PadOp([Id; 5]), // input, padding_value, edge_padding_low,
                                                        // edge_padding_high, interior_padding
@@ -50,22 +51,26 @@ define_language! {
       // Complete pain, has arity 12
       "ScatterOp"          = ScatterOp([Id; 4]), // input, scatter_indices, updates, dimension_numbers
       "ReturnOp"           = ReturnOp([Id; 1]),
-      "BlackBox"           = BlackBox([Id; 3]),  // id, args, captured values (last two should be vecs) 
+      "BlackBox"           = BlackBox([Id; 3]),  // id, args, captured values (last two should be vecs)
       "Vec"                = Vec(Vec<Id>),
       "Index"              = Index([Id; 2]),   // index, input. for indexing into ops with multiple result Values.
       // SHORTHANDS (not 1:1 with stablehlo)
       //
+      // Let axis' = max(0, len(shape(input)) - len(shape(orig_1))) + axis.
+      //
+      // TODO: We might need a symmetric "left leaning" version for a different rewrite
+      //
       // (SSplit0 input axis orig_0) means:
-      // split input on axis dimension, taking the left component. The split point is such that the result
-      // has the same shape as orig_0 on axis.
+      // split input on axis' dimension, taking the left component.
+      // The split point is such that the result has the same shape as orig_0 on axis'.
       //
       // (SSplit1 input axis orig_1) means:
-      // split input on axis dimension, taking the right component. The split point is such that the result
-      // has the same shape as orig_1 on axis.
-      //
+      // split input on axis' dimension, taking the right component.
+      // The split point is such that the result has the same shape as axis'.
+
       // This translates to a StableHLO SliceOp, with all the slices being [0..shape(input)[d]) in every
-      // dimension d except axis, and the axis slice being [0..shape(orig_0)[axis]) for SSplit0,
-      // and [shape(input)[axis] - shape(orig_1)[axis]..shape(input)[axis]) for SSplit1.
+      // dimension d except axis', and the axis slice being [0..shape(orig_0)[axis']) for SSplit0,
+      // and [shape(input)[axis'] - shape(orig_1)[axis']..shape(input)[axis']) for SSplit1.
       //
       // This allows embedding "splits" in syntactic rewrites (similarly to TASO) without keeping track of the
       // split tree, nor having a custom Applier to get the shape of inputs.
@@ -84,6 +89,10 @@ define_language! {
       // This allows certain rewrites with ConcatenateOps with axis larger than the rank of the operands, as
       // StableHLO doesn't have implicit casting.
       "MatchRank"         = MatchRank([Id; 2]),  // input, ref
+      // (InferReshape input) means:
+      // When this node is being merged into an eclass, use the shape of an existing eclass to create a ReshapeOp.
+      // Hence, this should only appear as the root node of RHS of a rewrite.
+      "InferReshape"      = InferReshape([Id; 1]),  // input
       // MISC
       Num(i64),
       Var(Symbol),
@@ -121,8 +130,11 @@ pub struct TensorData {
     // the operation.
     /// The list of results of this tensor
     pub tensors: Vec<ffi::Tensor>,
+    /// Is the root node InferReshape?
+    pub need_infer_shape: bool,
     /// The name string of this eclass if it is a Name type
     pub name: Option<&'static str>,
+
 }
 
 // Struct for storing information of a tensor. This is passed between functions
@@ -160,8 +172,22 @@ impl Analysis<Mdl> for TensorAnalysis {
     type Data = TensorData;
 
     fn merge(&self, to: &mut Self::Data, from: Self::Data) -> bool {
-        assert!(to.tensors == from.tensors, "{:?}{:?}", to, from);
-        false
+        assert!(!to.need_infer_shape || !from.need_infer_shape);
+        if !to.need_infer_shape && !from.need_infer_shape {
+            assert!(to.tensors == from.tensors, "to: {:?}, from: {:?}", to, from);
+            false
+        } else {
+            assert!(to.tensors.len() == from.tensors.len());
+            for i in 0..to.tensors.len() {
+                assert!(to.tensors[i].element_type == from.tensors[i].element_type);
+            }
+            if to.need_infer_shape {
+                *to = from;
+                true
+            } else {
+                false
+            }
+        }
     }
 
     fn make(egraph: &EGraph<Mdl, Self>, enode: &Mdl) -> Self::Data {
@@ -171,7 +197,7 @@ impl Analysis<Mdl> for TensorAnalysis {
         fn map_to_tensor(vec: Vec<i32>) -> ffi::Tensor {
             ffi::Tensor {
                 shape: vec.into_iter().map(|x| x as i64).collect(),
-                element_type: ffi::Type::i32, // Example: assuming i32 element type, modify as needed
+                element_type: ffi::Type::i32, 
             }
         }
 
@@ -187,10 +213,11 @@ impl Analysis<Mdl> for TensorAnalysis {
         match enode {
             Mdl::Num(_) | Mdl::Vec(_) => TensorData {
                 tensors: vec![ffi::Tensor {
-                    shape: vec![0],
+                    shape: vec![],
                     element_type: ffi::Type::i32,
                 }],
                 name: Some("num"),
+                need_infer_shape: false,
             },
             Mdl::Var(name) => {
                 let shape: Vec<i64> = name
@@ -217,6 +244,7 @@ impl Analysis<Mdl> for TensorAnalysis {
                         element_type,
                     }],
                     name: Some(name.as_str()),
+                    need_infer_shape: false,
                 }
             }
             Mdl::Input([node, _block_arg_number]) => x(node).clone(),
@@ -226,6 +254,7 @@ impl Analysis<Mdl> for TensorAnalysis {
                 TensorData {
                     tensors: vec![input.tensors[index].clone()],
                     name: None,
+                    need_infer_shape: false,
                 }
             }
             Mdl::BlackBox(inputs) => {
@@ -239,15 +268,29 @@ impl Analysis<Mdl> for TensorAnalysis {
                 TensorData {
                     tensors: shape_vec,
                     name: None,
+                    need_infer_shape: false,
                 }
             }
             Mdl::ReturnOp(_) => TensorData {
                 tensors: vec![],
                 name: None,
+                need_infer_shape: false,
+            },
+            Mdl::InferReshape([input]) => {
+                let input = x(input);
+                TensorData {
+                    tensors: vec![ffi::Tensor {
+                        shape: vec![],
+                        element_type: input.tensors[0].element_type,
+                    }],
+                    name: None,
+                    need_infer_shape: true,
+                }
             },
             x => TensorData {
                 tensors: create_stablehlo_op(egraph, x, ffi::get_shape),
                 name: None,
+                need_infer_shape: false,
             },
         }
     }
diff --git a/src/enzyme_ad/jax/deps/tensat/src/optimize.rs b/src/enzyme_ad/jax/deps/tensat/src/optimize.rs
index 8d5de8c9..2aff185d 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/optimize.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/optimize.rs
@@ -54,7 +54,10 @@ impl CostModel {
             | Mdl::Vec(_)
             | Mdl::BlackBox(_)
             | Mdl::Index(_)
-            | Mdl::ReturnOp(_) => 0.0,
+            | Mdl::ReturnOp(_)
+            // InferReshape at most adds a ReshapeOp without changing the underlying data, which should be free
+            // TODO: Gather all the "zero cost ops" assumptions together (the other bits currently being in EqualitySaturation.cpp)
+            | Mdl::InferReshape(_) => 0.0,
             x => create_stablehlo_op(egraph, x, ffi::get_cost) as f32,
         }
     }
@@ -191,6 +194,7 @@ pub struct SolvedResults {
 /// - `node_picked`: hashmap storing which node is picked for each EClass ID
 /// - `eclass`: The EClass ID that we aim to construct as root
 /// - `added_memo`: Map from EClass ID to RecExpr ID. Storing the eclasses that were already added
+/// - `to_egraph`: Map from RecExpr ID to EClass ID
 /// - `egraph`: E-graph of interest
 /// - `expr`: the RecExpr storing the optimized graph, it is constructed within this function
 ///
@@ -201,6 +205,7 @@ pub fn construct_best_rec(
     node_picked: &HashMap<Id, Mdl>,
     eclass: Id,
     added_memo: &mut HashMap<Id, Id>,
+    to_egraph: &mut HashMap<Id, Id>,
     egraph: &EGraph<Mdl, TensorAnalysis>,
     expr: &mut RecExpr<Mdl>,
 ) -> Id {
@@ -210,10 +215,11 @@ pub fn construct_best_rec(
         Some(id_expr) => *id_expr,
         None => {
             let node = node_picked.get(&id).unwrap().clone().map_children(|child| {
-                construct_best_rec(node_picked, child, added_memo, egraph, expr)
+                construct_best_rec(node_picked, child, added_memo, to_egraph, egraph, expr)
             });
             let id_expr = expr.add(node);
             assert!(added_memo.insert(id, id_expr).is_none());
+            assert!(to_egraph.insert(id_expr, id).is_none());
             id_expr
         }
     }
diff --git a/src/enzyme_ad/jax/deps/tensat/src/rewrites.rs b/src/enzyme_ad/jax/deps/tensat/src/rewrites.rs
index 898a6cc0..9e0d87c7 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/rewrites.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/rewrites.rs
@@ -102,7 +102,13 @@ pub fn rules<A: Analysis<Mdl>>() -> Vec<Rewrite<Mdl, A>> { vec![
         rw!("-concatenation-and-pooling-2"     ;"(poolmax ?kx ?ky ?sx ?sy ?p (concat 1 ?x ?y))"                                     => "(concat 1 (poolmax ?kx ?ky ?sx ?sy ?p ?x) (poolmax ?kx ?ky ?sx ?sy ?p ?y)) "               ),
 ]}
 
-fn add_to_rule_vec(rule_vec: &mut Vec<Rewrite<Mdl, TensorAnalysis>>, filter_after: bool, rule_name: String, lhs: &str, rhs: &str) {
+fn add_to_rule_vec(
+    rule_vec: &mut Vec<Rewrite<Mdl, TensorAnalysis>>,
+    filter_after: bool,
+    rule_name: String,
+    lhs: &str,
+    rhs: &str,
+) {
     let lhs: Pattern<Mdl> = lhs.parse().unwrap();
     let rhs: Pattern<Mdl> = rhs.parse().unwrap();
     rule_vec.push(rw!(rule_name; { lhs.clone() } => { CheckApply {
@@ -117,12 +123,30 @@ pub fn rules_from_str(rs: Vec<&str>, filter_after: bool) -> Vec<Rewrite<Mdl, Ten
     for (pos, rule) in rs.iter().enumerate() {
         if (rule.contains("<=>")) {
             let eqn: Vec<&str> = rule.split("<=>").collect();
-            add_to_rule_vec(&mut rule_vec, filter_after, format!("rule{}_l", pos), eqn[0], eqn[1]);
-            add_to_rule_vec(&mut rule_vec, filter_after, format!("rule{}_r", pos), eqn[1], eqn[0]);
+            add_to_rule_vec(
+                &mut rule_vec,
+                filter_after,
+                format!("rule{}_l", pos),
+                eqn[0],
+                eqn[1],
+            );
+            add_to_rule_vec(
+                &mut rule_vec,
+                filter_after,
+                format!("rule{}_r", pos),
+                eqn[1],
+                eqn[0],
+            );
         } else {
             assert!(rule.contains("=>"));
             let eqn: Vec<&str> = rule.split("=>").collect();
-            add_to_rule_vec(&mut rule_vec, filter_after, format!("rule{}", pos), eqn[0], eqn[1]);
+            add_to_rule_vec(
+                &mut rule_vec,
+                filter_after,
+                format!("rule{}", pos),
+                eqn[0],
+                eqn[1],
+            );
         }
     }
     rule_vec
@@ -152,6 +176,27 @@ pub static PRE_DEFINED_MULTI: &[&str] = &[
 
 // TODO: We should really clean these. Just dirty hacks for now to test out conditional rewrites
 
+pub fn get_matrix_option(
+    egraph: &EGraph<Mdl, TensorAnalysis>,
+    eclass: &EClass<Mdl, TensorData>,
+) -> Option<ffi::Matrix> {
+    // First, we get an optional Vec of Vec<Id> from the eclass
+    get_vec_option(eclass)
+        .map(|outer_vec| {
+            // Iterate over each Vec<Id> in the outer_vec to convert to a Vector
+            outer_vec
+                .iter()
+                .map(|&id| {
+                    // For each inner vector, get its corresponding vector of numbers
+                    get_vec_of_nums_option(egraph, &egraph[id])
+                        .map(|ffi_vec| ffi::Vector { vec: ffi_vec.vec })
+                })
+                .collect::<Option<Vec<ffi::Vector>>>()
+        })
+        // If all the vectors were successfully retrieved, wrap them in a Matrix
+        .and_then(|opt_mat| opt_mat.map(|mat| ffi::Matrix { mat }))
+}
+
 pub fn get_vec_of_nums(
     egraph: &EGraph<Mdl, TensorAnalysis>,
     eclass: &EClass<Mdl, TensorData>,
diff --git a/src/enzyme_ad/jax/enzyme_call.cc b/src/enzyme_ad/jax/enzyme_call.cc
index 88452c8f..a7f4fa94 100644
--- a/src/enzyme_ad/jax/enzyme_call.cc
+++ b/src/enzyme_ad/jax/enzyme_call.cc
@@ -290,6 +290,9 @@ class CpuKernel {
           case xla::PrimitiveType::F64:
             ty = "double";
             break;
+          case xla::PrimitiveType::PRED:
+            ty = "bool";
+            break;
           default: {
             std::string err;
             llvm::raw_string_ostream ess(err);
@@ -844,8 +847,10 @@ class CpuKernel {
 
     auto mod = GetLLVMFromJob("/enzyme_call/source.cpp", ss.str(), /*cpp*/ true,
                               pyargv_strs, llvm_ctx.get(), std::move(linkMod));
-    if (!mod)
+    if (!mod) {
+      llvm::errs() << "Source:\n" << ss.str() << "\n";
       throw pybind11::value_error("failed to compile C++");
+    }
     return std::make_tuple(std::move(mod), std::move(llvm_ctx), out_off,
                            tmpBuf);
   }
diff --git a/src/enzyme_ad/jax/primitives.py b/src/enzyme_ad/jax/primitives.py
index d40b2285..1adc03c5 100644
--- a/src/enzyme_ad/jax/primitives.py
+++ b/src/enzyme_ad/jax/primitives.py
@@ -13,11 +13,12 @@
 from jaxlib.mlir import ir
 from jaxlib.mlir.dialects import stablehlo, func
 from jax.lib import xla_client
-
 import jax.numpy as jnp
 
 from . import enzyme_call
 
+from .utils import default_nowheel_resource, default_linux_cflags
+
 LANG_CPP = enzyme_call.Language.CPP
 LANG_LLVM = enzyme_call.Language.LLVM
 LANG_MHLO = enzyme_call.Language.MHLO
@@ -315,7 +316,6 @@ def hlo_opts():
 exp_simplify<16>;
 slice_simplify<16>;
 convert_simplify<16>;
-reshape_simplify<16>;
 dynamic_slice_to_static<16>;
 dynamic_update_slice_elim<16>;
 concat_to_broadcast<16>;
@@ -428,7 +428,8 @@ def resource_dir():
 
     dn = os.path.dirname(enzyme_call.__file__)
     if os.getenv("ENZYME_BAZEL_NOWHEEL", None) is None:
-        res = os.path.join(
+        res = default_nowheel_resource(dn)
+        os.path.join(
             dn, "..", "..", "..", "external", "llvm-project", "clang", "staging"
         )
     else:
@@ -455,7 +456,7 @@ def cflags():
             "-fgnuc-version=4.2.1",
         )
     else:
-        res = ()
+        res = default_linux_cflags()
         if os.getenv("ENABLE_GDBLISTENER") is not None:
             res = res + (
                 "-debug-info-kind=standalone",
@@ -1183,7 +1184,6 @@ def cpp_call(
 
 
 def enzyme_jvp(arg_primals, arg_tangents, **kwargs):
-
     # TODO propagate activity info rather than make_zero
     def make_zero(tan, prim):
         return lax.zeros_like_array(prim) if type(tan) is ad.Zero else tan
@@ -1432,7 +1432,6 @@ def primal_partial_eval(trace, *args, **kwargs):
 def enzyme_vjp(shadow_rets, *prim_args, **kwargs):
     pipeline_options = kwargs["pipeline_options"]
     if pipeline_options.mlir_ad() and kwargs["lang"] == LANG_MHLO:
-
         passes = pipeline_options.pass_pipeline()
         start = passes.rindex("enzyme-wrap{")
         prev_passes = passes[:start]
diff --git a/src/enzyme_ad/jax/utils.py b/src/enzyme_ad/jax/utils.py
new file mode 100644
index 00000000..9e1cbe3c
--- /dev/null
+++ b/src/enzyme_ad/jax/utils.py
@@ -0,0 +1,12 @@
+import os
+from .primitives import *
+
+
+def default_nowheel_resource(dn):
+    return os.path.join(
+        dn, "..", "..", "..", "external", "llvm-project", "clang", "staging"
+    )
+
+
+def default_linux_cflags():
+    return ()
diff --git a/src/enzyme_ad/tools/enzymexlamlir-tblgen.cpp b/src/enzyme_ad/tools/enzymexlamlir-tblgen.cpp
index 906f9c02..b58ee08a 100644
--- a/src/enzyme_ad/tools/enzymexlamlir-tblgen.cpp
+++ b/src/enzyme_ad/tools/enzymexlamlir-tblgen.cpp
@@ -29,14 +29,14 @@ static llvm::cl::opt<ActionType> action(
     llvm::cl::values(clEnumValN(GenPopulatePatternsInterfaceImpl,
                                 "gen-populate-patterns-interface-impl", "")));
 
-llvm::StringRef getPopulateFunctionNameSuffix(llvm::Record *rec) {
+llvm::StringRef getPopulateFunctionNameSuffix(const llvm::Record *rec) {
   return rec->getName().ends_with("Op") ? rec->getName().drop_back(2)
                                         : rec->getName();
 }
 
 static bool emitPopulatePatterns(llvm::raw_ostream &os,
-                                 llvm::RecordKeeper &records) {
-  for (llvm::Record *rec :
+                                 const llvm::RecordKeeper &records) {
+  for (const llvm::Record *rec :
        records.getAllDerivedDefinitions("EnzymeHLOPatternOp")) {
     os << "void ";
     llvm::StringRef ns = rec->getValueAsString("cppNamespace");
@@ -54,8 +54,8 @@ static bool emitPopulatePatterns(llvm::raw_ostream &os,
 }
 
 static bool emitPopulatePatternsFuncDecls(llvm::raw_ostream &os,
-                                          llvm::RecordKeeper &records) {
-  for (llvm::Record *rec :
+                                          const llvm::RecordKeeper &records) {
+  for (const llvm::Record *rec :
        records.getAllDerivedDefinitions("EnzymeHLOPatternOp")) {
     llvm::StringRef ns = rec->getValueAsString("cppNamespace");
     if (ns.starts_with("::"))
@@ -70,8 +70,8 @@ static bool emitPopulatePatternsFuncDecls(llvm::raw_ostream &os,
 }
 
 static bool emitPopulatePatternsFuncDefs(llvm::raw_ostream &os,
-                                         llvm::RecordKeeper &records) {
-  for (llvm::Record *rec :
+                                         const llvm::RecordKeeper &records) {
+  for (const llvm::Record *rec :
        records.getAllDerivedDefinitions("EnzymeHLOPatternOp")) {
     os << "void ";
     llvm::StringRef ns = rec->getValueAsString("cppNamespace");
@@ -90,7 +90,8 @@ static bool emitPopulatePatternsFuncDefs(llvm::raw_ostream &os,
   return false;
 }
 
-static bool tablegenMain(llvm::raw_ostream &os, llvm::RecordKeeper &records) {
+static bool tablegenMain(llvm::raw_ostream &os,
+                         const llvm::RecordKeeper &records) {
   switch (action) {
   case GenPopulatePatternsFuncDecl:
     return emitPopulatePatternsFuncDecls(os, records);
diff --git a/workspace.bzl b/workspace.bzl
index 1d8fb464..a3cf23d6 100644
--- a/workspace.bzl
+++ b/workspace.bzl
@@ -1,7 +1,7 @@
-JAX_COMMIT = "9629bb8695e886d4bf0d25a454a0a157da7349f4"
+JAX_COMMIT = "81a31f6adf453b2afc39936e15c15d8ad327bf6e"
 JAX_SHA256 = ""
 
-ENZYME_COMMIT = "471c3c7ba8b02b7527003b02932e46710b160c1f"
+ENZYME_COMMIT = "192d9231aaa5415fa135ffb3568cc8eceef512d5"
 ENZYME_SHA256 = ""
 
 XLA_PATCHES = [


Test patch:
diff --git a/builddeps/test-requirements.txt b/builddeps/test-requirements.txt
index b27f584d..57fd36c0 100644
--- a/builddeps/test-requirements.txt
+++ b/builddeps/test-requirements.txt
@@ -2,7 +2,12 @@ absl-py
 jax
 numpy
 jaxlib
-jax-cuda12-plugin; sys_platform == 'linux'
+ortools
+https://github.com/wsmoses/jax-md/archive/45059b8f63dad0b5cb171feafff71b82162487e7.tar.gz
+# maxtext can't be installed concurrently, but installing it fixes
+# https://github.com/wsmoses/maxtext/archive/bc50722be7d89e4003bd830b80e4ac968be658eb.tar.gz
+jax[cuda12_pip]; sys_platform == 'linux'
 requests; sys_platform == 'linux'
--f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-libtpu-nightly == 0.1.dev20240729; sys_platform == 'linux' 
+nvidia-cuda-nvrtc-cu12; sys_platform == 'linux'
+# -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+# libtpu-nightly == 0.1.dev20240729; sys_platform == 'linux' 
diff --git a/test/BUILD b/test/BUILD
index 8e7fa522..3cafc54f 100644
--- a/test/BUILD
+++ b/test/BUILD
@@ -61,10 +61,11 @@ TEST_DEPS = [
         "//:enzyme_ad",
         "@pypi_jax//:pkg",
         "@pypi_absl_py//:pkg",
+        "@pypi_ortools//:pkg",
     ] + select({
         ":use_tpu": ["@pypi_libtpu_nightly//:pkg", "@pypi_requests//:pkg"],
-        "@bazel_tools//src/conditions:linux_x86_64": ["@pypi_jax_cuda12_plugin//:pkg"],
-        "//conditions:default": []
+        "@bazel_tools//src/conditions:linux_x86_64": ["@pypi_jax_cuda12_plugin//:pkg", "@pypi_nvidia_cuda_nvrtc_cu12//:pkg"],
+        "//conditions:default": [],
     })
 
 [
@@ -117,6 +118,27 @@ py_test(
     ],
     imports = ["."],
     deps = TEST_DEPS,
-    data = ["//src/enzyme_ad/jax:deps/tensat/converted.txt", "//src/enzyme_ad/jax:deps/tensat/extractor/extract.py"],
+    data = ["//src/enzyme_ad/jax:deps/tensat/converted.txt", "//src/enzyme_ad/jax:deps/tensat/converted_multi.txt", "//src/enzyme_ad/jax:deps/tensat/extractor/extract.py"],
     timeout='long'
 )
+
+py_test(
+    name = "jaxmd",
+    srcs = [
+        "jaxmd.py",
+        "test_utils.py",
+    ],
+    imports = ["."],
+    deps = TEST_DEPS + ["@pypi_jax_md//:pkg"],
+    timeout='long'
+)
+
+# py_test(
+#     name = "maxtext",
+#     srcs = [
+#         "maxtext.py",
+#     ],
+#     imports = ["."],
+#     deps = TEST_DEPS + ["@pypi_maxtext//:pkg"],
+#     timeout='long'
+# )
diff --git a/test/bench_vs_xla.py b/test/bench_vs_xla.py
index 748783ce..b0c1171e 100644
--- a/test/bench_vs_xla.py
+++ b/test/bench_vs_xla.py
@@ -258,5 +258,51 @@ def f(x, y):
                         )
 
 
+class ConstScatter(EnzymeJaxTest):
+    def setUp(self):
+        def forward(c_tau):
+            Q = c_tau
+            Q = Q.at[0].multiply(3)
+            chain = (Q,)
+            return (chain[0],)
+
+        self.fn = forward
+        self.name = "const_scatter"
+        self.count = 10
+
+        self.ins = [
+            jnp.array([2.7, 2.7, 2.7]),
+        ]
+        self.dins = [jnp.array([3.1, 3.1, 3.1])]
+        self.douts = self.dins
+        self.revfilter = lambda _: []
+        # No support for stablehlo.while atm
+        # self.revfilter = justjax
+        self.mlirad_rev = False
+
+
+class ScatterSum(EnzymeJaxTest):
+    def setUp(self):
+        def energy_fn(R, neighbor):
+            dR = R[neighbor[0]]
+            return jnp.sum(jnp.sin(dR))
+
+        nbrs = jnp.array([[2, 2]])
+
+        def forward(position):
+            e = jax.grad(energy_fn)(position, neighbor=nbrs)
+            return (e,)
+
+        self.fn = forward
+        self.name = "scatter_sum"
+
+        self.ins = [jnp.array([2.0, 4.0, 6.0, 8.0])]
+        self.dins = [jnp.array([2.7, 3.1, 5.9, 4.2])]
+        self.douts = []
+        self.revfilter = lambda _: []
+        # No support for stablehlo.scatter atm
+        self.mlirad_rev = False
+
+
 if __name__ == "__main__":
     absltest.main()
diff --git a/test/bert.py b/test/bert.py
new file mode 100644
index 00000000..e36f8e1b
--- /dev/null
+++ b/test/bert.py
@@ -0,0 +1,75 @@
+from transformers import FlaxBertModel, BertTokenizer
+from PIL import Image
+import jax
+import requests
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import absltest
+import jax.random
+from enzyme_ad.jax import JaXPipeline, hlo_opts
+from test_utils import *
+import llama
+
+# Define the pipelines as per your framework
+pipelines = [
+    ("JaX", None, CurBackends),
+    ("JaXPipe", JaXPipeline(), CurBackends),
+    (
+        "HLOOpt",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "canonicalize,cse,enzyme-hlo-opt,cse"
+        ),
+        CurBackends,
+    ),
+    ("PartOpt", JaXPipeline(llama.partialopt), CurBackends),
+    ("DefOpt", JaXPipeline(hlo_opts()), CurBackends),
+    (
+        "EqSat",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "equality-saturation-pass"
+        ),
+        CurBackends,
+    ),
+]
+
+# Load the tokenizer and model
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+model = FlaxBertModel.from_pretrained("bert-base-uncased")
+
+# Define the BERT forward pass
+def bert_forward(input_ids, attention_mask):
+    """ BERT forward pass """
+    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+    return outputs.last_hidden_state
+
+# Define the test class for BERT using absltest
+class BertTransformerTest(EnzymeJaxTest):
+    def setUp(self):
+        # Sample text input
+        text = "This is a sample text for BERT transformer test."
+        inputs = tokenizer(text, return_tensors="np")
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+
+        # Run the BERT forward pass with the processed input
+        self.hidden_state = bert_forward(input_ids, attention_mask)
+
+        # Setup for the test harness
+        self.fn = bert_forward
+        self.name = "bert"
+        self.count = 200
+        self.revprimal = False
+        self.AllPipelines = pipelines
+        self.AllBackends = CurBackends
+
+        # Input and output setup
+        self.ins = [input_ids, attention_mask]
+        self.dins = [input_ids, attention_mask]
+        self.douts = [self.hidden_state]
+        self.tol = 5e-5
+
+# Running the tests with absltest
+if __name__ == "__main__":
+    absltest.main()
diff --git a/test/jaxmd.py b/test/jaxmd.py
new file mode 100644
index 00000000..c9aa3617
--- /dev/null
+++ b/test/jaxmd.py
@@ -0,0 +1,258 @@
+from absl.testing import absltest
+import jax.numpy as jnp
+import jax.random
+import jax.lax
+import enzyme_ad.jax as enzyme_jax
+from enzyme_ad.jax import (
+    enzyme_jax_ir,
+    NewXLAPipeline,
+    OldXLAPipeline,
+    JaXPipeline,
+    hlo_opts,
+)
+import numpy as np
+import timeit
+from test_utils import *
+
+argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
+
+import jax.numpy as np
+import numpy as onp
+from jax import jit
+from jax import random
+from jax import lax
+
+jax.config.update("jax_enable_x64", True)
+
+from jax_md import space
+from jax_md import energy
+from jax_md import simulate
+from jax_md import quantity
+from jax_md import partition
+
+partialopt = (
+    "inline{default-pipeline=canonicalize max-iterations=4},"
+    + """canonicalize,cse,
+enzyme-hlo-generate-td{
+            patterns=compare_op_canon<16>;
+transpose_transpose<16>;
+broadcast_in_dim_op_canon<16>;
+convert_op_canon<16>;
+dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
+chained_dynamic_broadcast_in_dim_canonicalization<16>;
+dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
+noop_reduce_op_canon<16>;
+empty_reduce_op_canon<16>;
+dynamic_reshape_op_canon<16>;
+get_tuple_element_op_canon<16>;
+real_op_canon<16>;
+imag_op_canon<16>;
+get_dimension_size_op_canon<16>;
+gather_op_canon<16>;
+reshape_op_canon<16>;
+merge_consecutive_reshapes<16>;
+transpose_is_reshape<16>;
+zero_extent_tensor_canon<16>;
+reorder_elementwise_and_shape_op<16>;
+
+cse_broadcast_in_dim<16>;
+cse_slice<16>;
+cse_transpose<16>;
+cse_convert<16>;
+cse_pad<16>;
+cse_dot_general<16>;
+cse_reshape<16>;
+cse_mul<16>;
+cse_div<16>;
+cse_add<16>;
+cse_subtract<16>;
+cse_min<16>;
+cse_max<16>;
+cse_neg<16>;
+cse_concatenate<16>;
+
+concatenate_op_canon<16>(1024);
+select_op_canon<16>(1024);
+add_simplify<16>;
+sub_simplify<16>;
+and_simplify<16>;
+max_simplify<16>;
+min_simplify<16>;
+or_simplify<16>;
+negate_simplify<16>;
+mul_simplify<16>;
+div_simplify<16>;
+rem_simplify<16>;
+pow_simplify<16>;
+sqrt_simplify<16>;
+cos_simplify<16>;
+sin_simplify<16>;
+noop_slice<16>;
+const_prop_through_barrier<16>;
+slice_slice<16>;
+shift_right_logical_simplify<16>;
+pad_simplify<16>;
+negative_pad_to_slice<16>;
+tanh_simplify<16>;
+exp_simplify<16>;
+slice_simplify<16>;
+convert_simplify<16>;
+dynamic_slice_to_static<16>;
+dynamic_update_slice_elim<16>;
+concat_to_broadcast<16>;
+reduce_to_reshape<16>;
+broadcast_to_reshape<16>;
+gather_simplify<16>;
+iota_simplify<16>(1024);
+broadcast_in_dim_simplify<16>(1024);
+convert_concat<1>;
+dynamic_update_to_concat<1>;
+slice_of_dynamic_update<1>;
+slice_elementwise<1>;
+slice_pad<1>;
+dot_reshape_dot<1>;
+concat_const_prop<1>;
+concat_fuse<1>;
+pad_reshape_pad<1>;
+pad_pad<1>;
+concat_push_binop_add<1>;
+concat_push_binop_mul<1>;
+scatter_to_dynamic_update_slice<1>;
+reduce_concat<1>;
+slice_concat<1>;
+
+bin_broadcast_splat_add<1>;
+bin_broadcast_splat_subtract<1>;
+bin_broadcast_splat_div<1>;
+bin_broadcast_splat_mul<1>;
+slice_reshape<1>;
+
+dot_reshape_pad<1>;
+pad_dot_general<1>(1);
+broadcast_reduce<1>;
+            },
+            transform-interpreter,
+            enzyme-hlo-remove-transform,cse"""
+)
+
+pipelines = [
+    ("JaX  ", None, CurBackends),
+    ("JaXPipe", JaXPipeline(), CurBackends),
+    (
+        "HLOOpt",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "canonicalize,cse,enzyme-hlo-opt,cse"
+        ),
+        CurBackends,
+    ),
+    ("PartOpt", JaXPipeline(partialopt), CurBackends),
+    ("DefOpt", JaXPipeline(hlo_opts()), CurBackends),
+]
+
+
+class JAXMD(EnzymeJaxTest):
+    def setUp(self):
+        lattice_constant = 1.37820
+        N_rep = 40
+        box_size = N_rep * lattice_constant
+        # Using float32 for positions / velocities, but float64 for reductions.
+        dtype = np.float32
+
+        # Specify the format of the neighbor list.
+        # Options are Dense, Sparse, or OrderedSparse.
+        format = partition.OrderedSparse
+
+        displacement, shift = space.periodic(box_size)
+
+        R = []
+        for i in range(N_rep):
+            for j in range(N_rep):
+                for k in range(N_rep):
+                    R += [[i, j, k]]
+        R = np.array(R, dtype=dtype) * lattice_constant
+
+        N = R.shape[0]
+        phi = N / (lattice_constant * N_rep) ** 3
+        print(f"Created a system of {N} LJ particles with number density {phi:.3f}")
+
+        neighbor_fn, energy_fn = energy.lennard_jones_neighbor_list(
+            displacement, box_size, r_cutoff=3.0, dr_threshold=1.0, format=format
+        )
+
+        init, apply = simulate.nvt_nose_hoover(energy_fn, shift, 5e-3, kT=1.2)
+
+        key = random.PRNGKey(0)
+
+        # We pick an "extra capacity" to ensure ahead of time that the neighbor
+        # list will have enough capacity. Since sparse neighbor lists are more
+        # robust to changes in the number of particles, in this case we only
+        # need to actually add more capacity for dense neighbor lists.
+        if format is partition.Dense:
+            nbrs = neighbor_fn.allocate(R, extra_capacity=55)
+        else:
+            nbrs = neighbor_fn.allocate(R)
+
+        state = init(key, R, neighbor=nbrs)
+
+        def step(i, state_and_nbrs):
+            state, nbrs = state_and_nbrs
+            nbrs = nbrs.update(state.position)
+            return apply(state, neighbor=nbrs), nbrs
+
+        iters = 10
+        degrees_of_freedom = state.chain.degrees_of_freedom
+
+        def forward(
+            position, momentum, force, mass, c_position, c_momentum, c_mass, c_tau, c_KE
+        ):
+            chain = simulate.NoseHooverChain(
+                c_position, c_momentum, c_mass, c_tau, c_KE, degrees_of_freedom
+            )
+            state = simulate.NVTNoseHooverState(position, momentum, force, mass, chain)
+            # new_state, new_nbrs = lax.fori_loop(0, iters, step, (state, nbrs))
+            new_state, new_nbrs = step(0, (state, nbrs))
+            return (
+                new_state.position,
+                new_state.momentum,
+                new_state.force,
+                new_state.mass,
+                new_state.chain.position,
+                new_state.chain.momentum,
+                new_state.chain.mass,
+                new_state.chain.tau,
+                new_state.chain.kinetic_energy,
+            )
+
+        self.fn = forward
+        self.name = "jaxmd40"
+        self.count = 10
+        # self.revprimal = False
+        # self.AllPipelines = pipelines
+        # self.AllBackends = CurBackends
+
+        self.ins = [
+            state.position,
+            state.momentum,
+            state.force,
+            state.mass,
+            state.chain.position,
+            state.chain.momentum,
+            state.chain.mass,
+            state.chain.tau,
+            state.chain.kinetic_energy,
+        ]
+        # for i, v in enumerate(self.ins):
+        #    print("i=", i, v)
+        self.dins = [x.copy() for x in self.ins]
+        self.douts = [x.copy() for x in self.ins]
+        self.AllPipelines = pipelines
+        # No support for stablehlo.while atm
+        # self.revfilter = justjax
+        self.mlirad_rev = False
+
+        self.tol = 5e-4
+
+
+if __name__ == "__main__":
+    absltest.main()
diff --git a/test/keras_resnet.py b/test/keras_resnet.py
new file mode 100644
index 00000000..4196be07
--- /dev/null
+++ b/test/keras_resnet.py
@@ -0,0 +1,41 @@
+import numpy as np
+from keras.applications import ResNet50
+from keras.models import Model
+from keras.layers import Dense, GlobalAveragePooling2D
+from keras.optimizers import SGD
+from keras.datasets import mnist
+from keras.utils import to_categorical
+
+# Load the MNIST dataset
+(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
+
+# Convert to 3 channels by repeating the grayscale image
+train_images = np.stack((train_images,) * 3, axis=-1)
+test_images = np.stack((test_images,) * 3, axis=-1)
+
+# Resize images to 224x224
+train_images = np.resize(train_images, (train_images.shape[0], 224, 224, 3))
+test_images = np.resize(test_images, (test_images.shape[0], 224, 224, 3))
+
+# Normalize pixel values to be between 0 and 1
+train_images = train_images.astype('float32') / 255.0
+test_images = test_images.astype('float32') / 255.0
+train_labels = to_categorical(train_labels, num_classes=10)
+test_labels = to_categorical(test_labels, num_classes=10)
+
+base_model = ResNet50(weights=None, include_top=False, input_shape=(224, 224, 3))
+x = base_model.output
+x = GlobalAveragePooling2D()(x)
+
+x = Dense(1024, activation='relu')(x)
+
+predictions = Dense(10, activation='softmax')(x)
+
+model = Model(inputs=base_model.input, outputs=predictions)
+
+# Freeze all convolutional ResNet50 layers
+for layer in base_model.layers:
+    layer.trainable = False
+
+model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])
+model.fit(train_images, train_labels, epochs=10, batch_size=32, validation_data=(test_images, test_labels))
diff --git a/test/lit_tests/batchtests/transpose.mlir b/test/lit_tests/batchtests/transpose.mlir
new file mode 100644
index 00000000..a396476c
--- /dev/null
+++ b/test/lit_tests/batchtests/transpose.mlir
@@ -0,0 +1,15 @@
+// RUN: enzymexlamlir-opt %s --enzyme-batch | FileCheck %s
+  
+func.func private @relu_broadcast_scalar(%arg0: tensor<3x4xf64>) -> (tensor<4x3xf64>) {
+    %1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<3x4xf64>) -> tensor<4x3xf64>
+    return %1 : tensor<4x3xf64>
+  }
+  func.func @main(%arg0: tensor<2x5x3x4xf64>) -> (tensor<2x5x4x3xf64>) {
+    %1 = enzyme.batch @relu_broadcast_scalar(%arg0) {batch_shape = array<i64: 2, 5>} : (tensor<2x5x3x4xf64>) -> (tensor<2x5x4x3xf64>)
+    return %1 : tensor<2x5x4x3xf64>
+  }
+
+// CHECK:  func.func private @batched_relu_broadcast_scalar(%arg0: tensor<2x5x3x4xf64>) -> tensor<2x5x4x3xf64> {
+// CHECK-NEXT:    %0 = stablehlo.transpose %arg0, dims = [0, 1, 3, 2] : (tensor<2x5x3x4xf64>) -> tensor<2x5x4x3xf64>
+// CHECK-NEXT:    return %0 : tensor<2x5x4x3xf64>
+// CHECK-NEXT:  }
diff --git a/test/lit_tests/diffrules/stablehlo/gather.mlir b/test/lit_tests/diffrules/stablehlo/gather.mlir
new file mode 100644
index 00000000..3ba57f26
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/gather.mlir
@@ -0,0 +1,15 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// TODO: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func @main(%arg0 : tensor<64x3xf32>, %7 : tensor<45x1xi32>) -> tensor<45x3xf32> {
+    %8 = "stablehlo.gather"(%arg0, %7) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<64x3xf32>, tensor<45x1xi32>) -> tensor<45x3xf32>
+    return %8 : tensor<45x3xf32>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<64x3xf32>, %arg1: tensor<64x3xf32>, %arg2: tensor<45x1xi32>, %arg3: tensor<45x1xi32>) -> (tensor<45x3xf32>, tensor<45x3xf32>) {
+// FORWARD-NEXT:    %0 = "stablehlo.gather"(%arg1, %arg2) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<64x3xf32>, tensor<45x1xi32>) -> tensor<45x3xf32>
+// FORWARD-NEXT:    %1 = "stablehlo.gather"(%arg0, %arg2) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<64x3xf32>, tensor<45x1xi32>) -> tensor<45x3xf32>
+// FORWARD-NEXT:    return %1, %0 : tensor<45x3xf32>, tensor<45x3xf32>
+// FORWARD-NEXT:  }
diff --git a/test/lit_tests/diffrules/stablehlo/insort.mlir b/test/lit_tests/diffrules/stablehlo/insort.mlir
new file mode 100644
index 00000000..bd9790f1
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/insort.mlir
@@ -0,0 +1,29 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+func.func @main(%x : tensor<64000xi64>) -> tensor<64000xi64> {
+  %i = stablehlo.iota dim = 0 : tensor<64000xi64>
+  %r:2 = "stablehlo.sort"(%x, %i) <{dimension = 0 : i64, is_stable = true}> ({
+    ^bb0(%arg170: tensor<i64>, %arg171: tensor<i64>, %arg172: tensor<i64>, %arg173: tensor<i64>):
+      %1096 = "stablehlo.compare"(%arg170, %arg171) <{compare_type = #stablehlo<comparison_type SIGNED>, comparison_direction = #stablehlo<comparison_direction LT>}> : (tensor<i64>, tensor<i64>) -> tensor<i1>
+      "stablehlo.return"(%1096) : (tensor<i1>) -> ()
+    }) : (tensor<64000xi64>, tensor<64000xi64>) -> (tensor<64000xi64>, tensor<64000xi64>)
+
+  func.return %r#1 : tensor<64000xi64>
+}
+
+// FORWARD:   func.func @main(%arg0: tensor<64000xi64>, %arg1: tensor<64000xi64>) -> (tensor<64000xi64>, tensor<64000xi64>) {
+// FORWARD-NEXT:     %0 = stablehlo.iota dim = 0 : tensor<64000xi64>
+// FORWARD-NEXT:     %1:2 = "stablehlo.sort"(%arg0, %0) <{dimension = 0 : i64, is_stable = true}> ({
+// FORWARD-NEXT:     ^bb0(%arg2: tensor<i64>, %arg3: tensor<i64>, %arg4: tensor<i64>, %arg5: tensor<i64>):
+// FORWARD-NEXT:       %2 = stablehlo.compare  LT, %arg2, %arg3,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
+// FORWARD-NEXT:       stablehlo.return %2 : tensor<i1>
+// FORWARD-NEXT:     }) : (tensor<64000xi64>, tensor<64000xi64>) -> (tensor<64000xi64>, tensor<64000xi64>)
+// FORWARD-NEXT:     %cst = arith.constant dense<0> : tensor<64000xi64>
+// FORWARD-NEXT:     return %1#1, %cst : tensor<64000xi64>, tensor<64000xi64>
+// FORWARD-NEXT:   }
+
+// REVERSE:   func.func @main(%arg0: tensor<64000xi64>, %arg1: tensor<64000xi64>) -> tensor<64000xi64> {
+// REVERSE-NEXT:     %cst = arith.constant dense<0> : tensor<64000xi64>
+// REVERSE-NEXT:     return %cst : tensor<64000xi64>
+// REVERSE-NEXT:   }
diff --git a/test/lit_tests/diffrules/stablehlo/reduce.mlir b/test/lit_tests/diffrules/stablehlo/reduce.mlir
new file mode 100644
index 00000000..045ad746
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/reduce.mlir
@@ -0,0 +1,25 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+func.func @main(%x : tensor<3680000x3xf32>) -> tensor<3680000xf32> {
+  %cst_76 = stablehlo.constant dense<2.7> : tensor<f32>
+  %res = stablehlo.reduce(%x init: %cst_76) applies stablehlo.add across dimensions = [1] : (tensor<3680000x3xf32>, tensor<f32>) -> tensor<3680000xf32>
+  func.return %res : tensor<3680000xf32>
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<3680000x3xf32>, %arg1: tensor<3680000x3xf32>) -> (tensor<3680000xf32>, tensor<3680000xf32>) {
+// FORWARD-NEXT:    %cst = stablehlo.constant dense<2.700000e+00> : tensor<f32>
+// FORWARD-NEXT:    %cst_0 = arith.constant dense<0.000000e+00> : tensor<f32>
+// FORWARD-NEXT:    %0 = stablehlo.reduce(%arg1 init: %cst_0) applies stablehlo.add across dimensions = [1] : (tensor<3680000x3xf32>, tensor<f32>) -> tensor<3680000xf32>
+// FORWARD-NEXT:    %1 = stablehlo.reduce(%arg0 init: %cst) applies stablehlo.add across dimensions = [1] : (tensor<3680000x3xf32>, tensor<f32>) -> tensor<3680000xf32>
+// FORWARD-NEXT:    return %1, %0 : tensor<3680000xf32>, tensor<3680000xf32>
+// FORWARD-NEXT:  }
+
+// REVERSE:  func.func @main(%arg0: tensor<3680000x3xf32>, %arg1: tensor<3680000xf32>) -> tensor<3680000x3xf32> {
+// REVERSE-NEXT:    %cst = arith.constant dense<0.000000e+00> : tensor<3680000xf32>
+// REVERSE-NEXT:    %cst_0 = arith.constant dense<0.000000e+00> : tensor<3680000x3xf32>
+// REVERSE-NEXT:    %0 = arith.addf %arg1, %cst : tensor<3680000xf32>
+// REVERSE-NEXT:    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<3680000xf32>) -> tensor<3680000x3xf32>
+// REVERSE-NEXT:    %2 = arith.addf %1, %cst_0 : tensor<3680000x3xf32>
+// REVERSE-NEXT:    return %2 : tensor<3680000x3xf32>
+// REVERSE-NEXT:  }
diff --git a/test/lit_tests/diffrules/stablehlo/reduceactivity.mlir b/test/lit_tests/diffrules/stablehlo/reduceactivity.mlir
new file mode 100644
index 00000000..773ba65e
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/reduceactivity.mlir
@@ -0,0 +1,24 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func public @main(%arg0: tensor<64000x3xf32>) -> (tensor<i32>) {
+    %c_39 = stablehlo.constant dense<64000> : tensor<i32>
+    %cst_76 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+    %11 = stablehlo.reduce(%arg0 init: %cst_76) applies stablehlo.add across dimensions = [1] : (tensor<64000x3xf32>, tensor<f32>) -> tensor<64000xf32>
+    return %c_39 : tensor<i32>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<64000x3xf32>, %arg1: tensor<64000x3xf32>) -> (tensor<i32>, tensor<i32>) {
+// FORWARD-NEXT:    %c = stablehlo.constant dense<64000> : tensor<i32>
+// FORWARD-NEXT:    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+// FORWARD-NEXT:    %0 = stablehlo.reduce(%arg0 init: %cst) applies stablehlo.add across dimensions = [1] : (tensor<64000x3xf32>, tensor<f32>) -> tensor<64000xf32>
+// FORWARD-NEXT:    %cst_0 = arith.constant dense<0> : tensor<i32>
+// FORWARD-NEXT:    return %c, %cst_0 : tensor<i32>, tensor<i32>
+// FORWARD-NEXT:  }
+
+// REVERSE:  func.func @main(%arg0: tensor<64000x3xf32>, %arg1: tensor<i32>) -> tensor<64000x3xf32> {
+// REVERSE-NEXT:    %cst = arith.constant dense<0.000000e+00> : tensor<64000x3xf32>
+// REVERSE-NEXT:    return %cst : tensor<64000x3xf32>
+// REVERSE-NEXT:  }
diff --git a/test/lit_tests/diffrules/stablehlo/scatter.mlir b/test/lit_tests/diffrules/stablehlo/scatter.mlir
new file mode 100644
index 00000000..e4cf2ef6
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/scatter.mlir
@@ -0,0 +1,22 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// TODO: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func @main(%prev : tensor<1x1x8192x16x256xbf16>, %update : tensor<1x1x2048x16x256xbf16>) -> tensor<1x1x8192x16x256xbf16> {
+    %163 = stablehlo.constant dense<0> : tensor<1xi32>
+    %524 = "stablehlo.scatter"(%prev, %163, %update) ({
+    ^bb0(%arg113: tensor<bf16>, %arg114: tensor<bf16>):
+      stablehlo.return %arg114 : tensor<bf16>
+    }) {indices_are_sorted = true, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0, 1, 2, 3, 4], scatter_dims_to_operand_dims = [2]>, unique_indices = true} : (tensor<1x1x8192x16x256xbf16>, tensor<1xi32>, tensor<1x1x2048x16x256xbf16>) -> tensor<1x1x8192x16x256xbf16>
+    return %524 : tensor<1x1x8192x16x256xbf16>
+  }
+}
+
+// FORWARD: func.func @main(%arg0: tensor<1x1x8192x16x256xbf16>, %arg1: tensor<1x1x8192x16x256xbf16>, %arg2: tensor<1x1x2048x16x256xbf16>, %arg3: tensor<1x1x2048x16x256xbf16>) -> (tensor<1x1x8192x16x256xbf16>, tensor<1x1x8192x16x256xbf16>) {
+// FORWARD-NEXT:    %c = stablehlo.constant dense<0> : tensor<1xi32>
+// FORWARD-NEXT:    %0:2 = "stablehlo.scatter"(%arg0, %arg1, %c, %arg2, %arg3) <{indices_are_sorted = true, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0, 1, 2, 3, 4], scatter_dims_to_operand_dims = [2]>, unique_indices = true}> ({
+// FORWARD-NEXT:    ^bb0(%arg4: tensor<bf16>, %arg5: tensor<bf16>, %arg6: tensor<bf16>, %arg7: tensor<bf16>):
+// FORWARD-NEXT:      stablehlo.return %arg6, %arg7 : tensor<bf16>, tensor<bf16>
+// FORWARD-NEXT:    }) : (tensor<1x1x8192x16x256xbf16>, tensor<1x1x8192x16x256xbf16>, tensor<1xi32>, tensor<1x1x2048x16x256xbf16>, tensor<1x1x2048x16x256xbf16>) -> (tensor<1x1x8192x16x256xbf16>, tensor<1x1x8192x16x256xbf16>)
+// FORWARD-NEXT:    return %0#0, %0#1 : tensor<1x1x8192x16x256xbf16>, tensor<1x1x8192x16x256xbf16>
+// FORWARD-NEXT:  }
diff --git a/test/lit_tests/diffrules/stablehlo/scatter1.mlir b/test/lit_tests/diffrules/stablehlo/scatter1.mlir
new file mode 100644
index 00000000..b87a1030
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/scatter1.mlir
@@ -0,0 +1,35 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// TODO: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+func.func @main(%x : tensor<5xf32>) -> tensor<5xf32> {
+  %c_30 = stablehlo.constant dense<0> : tensor<1xi32>
+  %cst_65 = stablehlo.constant dense<1.920000e+05> : tensor<f32>
+  %res = "stablehlo.scatter"(%x, %c_30, %cst_65) <{indices_are_sorted = true, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0]>, unique_indices = true}> ({
+    ^bb0(%arg8: tensor<f32>, %arg9: tensor<f32>):
+      %153 = stablehlo.multiply %arg8, %arg9 : tensor<f32>
+      stablehlo.return %153 : tensor<f32>
+  }) : (tensor<5xf32>, tensor<1xi32>, tensor<f32>) -> tensor<5xf32>
+  func.return %res : tensor<5xf32>
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<5xf32>, %arg1: tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>) {
+// FORWARD-NEXT:    %c = stablehlo.constant dense<0> : tensor<1xi32>
+// FORWARD-NEXT:    %cst = stablehlo.constant dense<1.920000e+05> : tensor<f32>
+// FORWARD-NEXT:    %cst_0 = arith.constant dense<0.000000e+00> : tensor<f32>
+// FORWARD-NEXT:    %0:2 = "stablehlo.scatter"(%arg0, %arg1, %c, %cst, %cst_0) <{indices_are_sorted = true, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0]>, unique_indices = true}> ({
+// FORWARD-NEXT:    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>, %arg4: tensor<f32>, %arg5: tensor<f32>):
+// FORWARD-NEXT:      %1 = stablehlo.multiply %arg3, %arg4 : tensor<f32>
+// FORWARD-NEXT:      %2 = stablehlo.multiply %arg2, %arg4 : tensor<f32>
+// FORWARD-NEXT:      stablehlo.return %2, %1 : tensor<f32>, tensor<f32>
+// FORWARD-NEXT:    }) : (tensor<5xf32>, tensor<5xf32>, tensor<1xi32>, tensor<f32>, tensor<f32>) -> (tensor<5xf32>, tensor<5xf32>)
+// FORWARD-NEXT:    return %0#0, %0#1 : tensor<5xf32>, tensor<5xf32>
+// FORWARD-NEXT:  }
+
+// REVERSE:  func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
+// REVERSE-NEXT:    %cst = arith.constant dense<0.000000e+00> : tensor<2xf32>
+// REVERSE-NEXT:    %0 = arith.addf %arg1, %cst : tensor<2xf32>
+// REVERSE-NEXT:    %1 = stablehlo.cosine %arg0 : tensor<2xf32>
+// REVERSE-NEXT:    %2 = stablehlo.multiply %0, %1 : tensor<2xf32>
+// REVERSE-NEXT:    %3 = arith.addf %2, %cst : tensor<2xf32>
+// REVERSE-NEXT:    return %3 : tensor<2xf32>
+// REVERSE-NEXT:  }
diff --git a/test/lit_tests/diffrules/stablehlo/scatteractivity.mlir b/test/lit_tests/diffrules/stablehlo/scatteractivity.mlir
new file mode 100644
index 00000000..ea9b3a6f
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/scatteractivity.mlir
@@ -0,0 +1,22 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// TODO: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_dup mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func public @main(%arg3: tensor<3xf32>, %arg6: tensor<5xf32>) -> (tensor<5xf32>) {
+    %c_0 = stablehlo.constant dense<[[3], [2], [1]]> : tensor<3x1xi32>
+    %33 = "stablehlo.scatter"(%arg6, %c_0, %arg3) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
+    ^bb0(%arg9: tensor<f32>, %arg10: tensor<f32>):
+      stablehlo.return %arg10 : tensor<f32>
+    }) : (tensor<5xf32>, tensor<3x1xi32>, tensor<3xf32>) -> tensor<5xf32>
+    return %33 : tensor<5xf32>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>, %arg2: tensor<5xf32>, %arg3: tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>) {
+// FORWARD-NEXT:    %c = stablehlo.constant dense<{{\[\[}}3], [2], [1{{\]\]}}> : tensor<3x1xi32>
+// FORWARD-NEXT:    %0:2 = "stablehlo.scatter"(%arg2, %arg3, %c, %arg0, %arg1) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
+// FORWARD-NEXT:    ^bb0(%arg4: tensor<f32>, %arg5: tensor<f32>, %arg6: tensor<f32>, %arg7: tensor<f32>):
+// FORWARD-NEXT:      stablehlo.return %arg6, %arg7 : tensor<f32>, tensor<f32>
+// FORWARD-NEXT:    }) : (tensor<5xf32>, tensor<5xf32>, tensor<3x1xi32>, tensor<3xf32>, tensor<3xf32>) -> (tensor<5xf32>, tensor<5xf32>)
+// FORWARD-NEXT:    return %0#0, %0#1 : tensor<5xf32>, tensor<5xf32>
+// FORWARD-NEXT:  }
diff --git a/test/lit_tests/diffrules/stablehlo/while.mlir b/test/lit_tests/diffrules/stablehlo/while.mlir
new file mode 100644
index 00000000..65cfb892
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/while.mlir
@@ -0,0 +1,35 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// TODO: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func public @main(%arg0: tensor<3xf64>) -> (tensor<3xf64>) {
+    %c = stablehlo.constant dense<0> : tensor<i64>
+    %c_10 = stablehlo.constant dense<10> : tensor<i64>
+    %c_1 = stablehlo.constant dense<1> : tensor<i64>
+    %0:2 = stablehlo.while(%iterArg = %c, %iterArg_0 = %arg0) : tensor<i64>, tensor<3xf64>
+     cond {
+      %1 = stablehlo.compare  LT, %iterArg, %c_10,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
+      stablehlo.return %1 : tensor<i1>
+    } do {
+      %2 = stablehlo.add %iterArg, %c_1 : tensor<i64>
+      stablehlo.return %2, %iterArg_0 : tensor<i64>, tensor<3xf64>
+    }
+    return %0#1 : tensor<3xf64>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<3xf64>, %arg1: tensor<3xf64>) -> (tensor<3xf64>, tensor<3xf64>) {
+// FORWARD-NEXT:    %c = stablehlo.constant dense<0> : tensor<i64>
+// FORWARD-NEXT:    %c_0 = stablehlo.constant dense<10> : tensor<i64>
+// FORWARD-NEXT:    %c_1 = stablehlo.constant dense<1> : tensor<i64>
+// FORWARD-NEXT:    %0:3 = stablehlo.while(%iterArg = %c, %iterArg_2 = %arg0, %iterArg_3 = %arg1) : tensor<i64>, tensor<3xf64>, tensor<3xf64>
+// FORWARD-NEXT:     cond {
+// FORWARD-NEXT:      %1 = stablehlo.compare  LT, %iterArg, %c_0,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
+// FORWARD-NEXT:      stablehlo.return %1 : tensor<i1>
+// FORWARD-NEXT:    } do {
+// FORWARD-NEXT:      %1 = stablehlo.add %iterArg, %c_1 : tensor<i64>
+// FORWARD-NEXT:      stablehlo.return %1, %iterArg_2, %iterArg_3 : tensor<i64>, tensor<3xf64>, tensor<3xf64>
+// FORWARD-NEXT:    }
+// FORWARD-NEXT:    return %0#1, %0#2 : tensor<3xf64>, tensor<3xf64>
+// FORWARD-NEXT:  }
+
diff --git a/test/lit_tests/diffrules/stablehlo/while2.mlir b/test/lit_tests/diffrules/stablehlo/while2.mlir
new file mode 100644
index 00000000..29df6277
--- /dev/null
+++ b/test/lit_tests/diffrules/stablehlo/while2.mlir
@@ -0,0 +1,36 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// TODO: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func public @main(%arg5: tensor<5xf32>, %1: tensor<f32>) -> (tensor<f32>) {
+    %c_2 = stablehlo.constant dense<1> : tensor<i64>
+    %c_3 = stablehlo.constant dense<0> : tensor<i64>
+    %2:3 = stablehlo.while(%iterArg_6 = %arg5, %iterArg_7 = %c_3, %iterArg_8 = %1) : tensor<5xf32>, tensor<i64>, tensor<f32>
+     cond {
+      %18 = stablehlo.compare  LT, %iterArg_7, %c_2,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
+      stablehlo.return %18 : tensor<i1>
+    } do {
+      %29 = stablehlo.dynamic_slice %iterArg_6, %c_3, sizes = [1] : (tensor<5xf32>, tensor<i64>) -> tensor<1xf32>
+      %30 = stablehlo.reshape %29 : (tensor<1xf32>) -> tensor<f32>
+      stablehlo.return %iterArg_6, %c_3, %30 : tensor<5xf32>, tensor<i64>, tensor<f32>
+    }
+    return %2#2 : tensor<f32>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<5xf32>, %arg1: tensor<5xf32>, %arg2: tensor<f32>, %arg3: tensor<f32>) -> (tensor<f32>, tensor<f32>)
+// FORWARD-NEXT:    %c = stablehlo.constant dense<1> : tensor<i64>
+// FORWARD-NEXT:    %c_0 = stablehlo.constant dense<0> : tensor<i64>
+// FORWARD-NEXT:    %0:5 = stablehlo.while(%iterArg = %arg0, %iterArg_1 = %arg1, %iterArg_2 = %c_0, %iterArg_3 = %arg2, %iterArg_4 = %arg3) : tensor<5xf32>, tensor<5xf32>, tensor<i64>, tensor<f32>, tensor<f32>
+// FORWARD-NEXT:     cond {
+// FORWARD-NEXT:      %1 = stablehlo.compare  LT, %iterArg_2, %c,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
+// FORWARD-NEXT:      stablehlo.return %1 : tensor<i1>
+// FORWARD-NEXT:    } do {
+// FORWARD-NEXT:      %1 = stablehlo.dynamic_slice %iterArg_1, %c_0, sizes = [1] : (tensor<5xf32>, tensor<i64>) -> tensor<1xf32>
+// FORWARD-NEXT:      %2 = stablehlo.dynamic_slice %iterArg, %c_0, sizes = [1] : (tensor<5xf32>, tensor<i64>) -> tensor<1xf32>
+// FORWARD-NEXT:      %3 = stablehlo.reshape %1 : (tensor<1xf32>) -> tensor<f32>
+// FORWARD-NEXT:      %4 = stablehlo.reshape %2 : (tensor<1xf32>) -> tensor<f32>
+// FORWARD-NEXT:      stablehlo.return %iterArg, %iterArg_1, %c_0, %4, %3 : tensor<5xf32>, tensor<5xf32>, tensor<i64>, tensor<f32>, tensor<f32>
+// FORWARD-NEXT:    }
+// FORWARD-NEXT:    return %0#3, %0#4 : tensor<f32>, tensor<f32>
+// FORWARD-NEXT:  }
diff --git a/test/lit_tests/nz.mlir b/test/lit_tests/nz.mlir
new file mode 100644
index 00000000..a50a3e99
--- /dev/null
+++ b/test/lit_tests/nz.mlir
@@ -0,0 +1,14 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+module {
+  func.func @wrapped() -> (tensor<1x3xf32> {mhlo.layout_mode = "default"}) {
+    %cst_0 = stablehlo.constant dense<[[0.000000e+00, 0.000000e+00, 0.000000e+00], [0.000000e+00, 0.000000e+00, 1.378200e+00], [0.000000e+00, 1.378200e+00, 0.000000e+00], [0.000000e+00, 1.378200e+00, 1.378200e+00], [1.378200e+00, 0.000000e+00, 0.000000e+00], [1.378200e+00, 0.000000e+00, 1.378200e+00], [1.378200e+00, 1.378200e+00, 0.000000e+00], [1.378200e+00, 1.378200e+00, 1.378200e+00]]> : tensor<8x3xf32>
+    %10 = stablehlo.slice %cst_0 [1:2, 0:3] : (tensor<8x3xf32>) -> tensor<1x3xf32>
+    return %10 : tensor<1x3xf32>
+  }
+}
+
+// CHECK:  func.func @wrapped() -> (tensor<1x3xf32> {mhlo.layout_mode = "default"}) {
+// CHECK-NEXT:    %cst = stablehlo.constant dense<{{\[\[}}0.000000e+00, 0.000000e+00, 1.378200e+00{{\]\]}}> : tensor<1x3xf32>
+// CHECK-NEXT:    return %cst : tensor<1x3xf32>
+// CHECK-NEXT:  }
diff --git a/test/llama-old.py b/test/llama-old.py
deleted file mode 100644
index 949edba8..00000000
--- a/test/llama-old.py
+++ /dev/null
@@ -1,829 +0,0 @@
-from absl.testing import absltest
-import jax.numpy as jnp
-import jax.random
-import jax.lax
-import enzyme_ad.jax as enzyme_jax
-import numpy as np
-import timeit
-
-argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
-
-massive_pipeline = enzyme_jax.JaXPipeline("""
-builtin.module(
-inline{default-pipeline=canonicalize max-iterations=4},
-canonicalize,cse,
-canonicalize,
-enzyme-hlo-generate-td{
-patterns=compare_op_canon<16>;
-transpose_transpose<16>;
-broadcast_in_dim_op_canon<16>;
-convert_op_canon<16>;
-dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
-chained_dynamic_broadcast_in_dim_canonicalization<16>;
-dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
-noop_reduce_op_canon<16>;
-empty_reduce_op_canon<16>;
-dynamic_reshape_op_canon<16>;
-get_tuple_element_op_canon<16>;
-real_op_canon<16>;
-imag_op_canon<16>;
-get_dimension_size_op_canon<16>;
-gather_op_canon<16>;
-reshape_op_canon<16>;
-merge_consecutive_reshapes<16>;
-transpose_is_reshape<16>;
-zero_extent_tensor_canon<16>;
-reorder_elementwise_and_shape_op<16>;
-
-cse_broadcast_in_dim<16>;
-cse_slice<16>;
-cse_transpose<16>;
-cse_convert<16>;
-cse_pad<16>;
-cse_dot_general<16>;
-cse_reshape<16>;
-cse_mul<16>;
-cse_div<16>;
-cse_add<16>;
-cse_subtract<16>;
-cse_min<16>;
-cse_max<16>;
-cse_neg<16>;
-cse_concatenate<16>;
-
-concatenate_op_canon<16>(1024);
-select_op_canon<16>(1024);
-add_simplify<16>;
-sub_simplify<16>;
-and_simplify<16>;
-max_simplify<16>;
-min_simplify<16>;
-or_simplify<16>;
-negate_simplify<16>;
-mul_simplify<16>;
-div_simplify<16>;
-rem_simplify<16>;
-pow_simplify<16>;
-sqrt_simplify<16>;
-cos_simplify<16>;
-sin_simplify<16>;
-noop_slice<16>;
-const_prop_through_barrier<16>;
-slice_slice<16>;
-shift_right_logical_simplify<16>;
-pad_simplify<16>;
-negative_pad_to_slice<16>;
-tanh_simplify<16>;
-exp_simplify<16>;
-slice_simplify<16>;
-convert_simplify<16>;
-reshape_simplify<16>;
-dynamic_slice_to_static<16>;
-dynamic_update_slice_elim<16>;
-concat_to_broadcast<16>;
-reduce_to_reshape<16>;
-broadcast_to_reshape<16>;
-gather_simplify<16>;
-iota_simplify<16>(1024);
-broadcast_in_dim_simplify<16>(1024);
-convert_concat<1>;
-dynamic_update_to_concat<1>;
-slice_of_dynamic_update<1>;
-slice_elementwise<1>;
-slice_pad<1>;
-dot_reshape_dot<1>;
-concat_const_prop<1>;
-concat_fuse<1>;
-pad_reshape_pad<1>;
-pad_pad<1>;
-concat_push_binop_add<1>;
-concat_push_binop_mul<1>;
-scatter_to_dynamic_update_slice<1>;
-reduce_concat<1>;
-slice_concat<1>;
-
-bin_broadcast_splat_add<1>;
-bin_broadcast_splat_subtract<1>;
-bin_broadcast_splat_div<1>;
-bin_broadcast_splat_mul<1>;
-slice_reshape<1>;
-
-dot_reshape_pad<1>;
-pad_dot_general<1>(1);
-broadcast_reduce<1>;},
-transform-interpreter,
-enzyme-hlo-remove-transform
-)""")
-
-massive_pipeline_with_eqsat = enzyme_jax.JaXPipeline("""
-inline{default-pipeline=canonicalize max-iterations=4},
-equality-saturation-pass,
-builtin.module(
-inline{default-pipeline=canonicalize max-iterations=4},
-canonicalize,cse,
-canonicalize,
-enzyme-hlo-generate-td{
-patterns=
-compare_op_canon<16>;
-broadcast_in_dim_op_canon<16>;
-convert_op_canon<16>;
-dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
-chained_dynamic_broadcast_in_dim_canonicalization<16>;
-dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
-noop_reduce_op_canon<16>;
-empty_reduce_op_canon<16>;
-dynamic_reshape_op_canon<16>;
-get_tuple_element_op_canon<16>;
-real_op_canon<16>;
-imag_op_canon<16>;
-get_dimension_size_op_canon<16>;
-gather_op_canon<16>;
-reshape_op_canon<16>;
-merge_consecutive_reshapes<16>;
-transpose_is_reshape<16>;
-zero_extent_tensor_canon<16>;
-reorder_elementwise_and_shape_op<16>;
-cse_broadcast_in_dim<16>;
-cse_slice<16>;
-cse_transpose<16>;
-cse_convert<16>;
-cse_pad<16>;
-cse_dot_general<16>;
-cse_reshape<16>;
-cse_mul<16>;
-cse_div<16>;
-cse_add<16>;
-cse_subtract<16>;
-cse_min<16>;
-cse_max<16>;
-cse_neg<16>;
-cse_concatenate<16>;
-concatenate_op_canon<16>(1024);
-select_op_canon<16>(1024);
-add_simplify<16>;
-sub_simplify<16>;
-and_simplify<16>;
-max_simplify<16>;
-min_simplify<16>;
-or_simplify<16>;
-negate_simplify<16>;
-mul_simplify<16>;
-div_simplify<16>;
-rem_simplify<16>;
-pow_simplify<16>;
-sqrt_simplify<16>;
-cos_simplify<16>;
-sin_simplify<16>;
-noop_slice<16>;
-const_prop_through_barrier<16>;
-slice_slice<16>;
-shift_right_logical_simplify<16>;
-pad_simplify<16>;
-negative_pad_to_slice<16>;
-tanh_simplify<16>;
-exp_simplify<16>;
-slice_simplify<16>;
-convert_simplify<16>;
-reshape_simplify<16>;
-dynamic_slice_to_static<16>;
-dynamic_update_slice_elim<16>;
-concat_to_broadcast<16>;
-reduce_to_reshape<16>;
-broadcast_to_reshape<16>;
-gather_simplify<16>;
-iota_simplify<16>(1024);
-broadcast_in_dim_simplify<16>(1024);
-convert_concat<1>;
-dynamic_update_to_concat<1>;
-slice_of_dynamic_update<1>;
-slice_elementwise<1>;
-slice_pad<1>;
-dot_reshape_dot<1>;
-concat_const_prop<1>;
-concat_fuse<1>;
-pad_reshape_pad<1>;
-pad_pad<1>;
-concat_push_binop_add<1>;
-concat_push_binop_mul<1>;
-scatter_to_dynamic_update_slice<1>;
-reduce_concat<1>;
-slice_concat<1>;
-bin_broadcast_splat_add<1>;
-bin_broadcast_splat_subtract<1>;
-bin_broadcast_splat_div<1>;
-bin_broadcast_splat_mul<1>;
-reshape_iota<16>;
-slice_reshape_slice<1>;
-dot_general_simplify<16>;
-transpose_simplify<16>;
-reshape_empty_broadcast<1>;
-add_pad_pad_to_concat<1>;
-broadcast_reshape<1>;
-slice_reshape_concat<1>;
-slice_reshape_elementwise<1>;
-slice_reshape_transpose<1>;
-slice_reshape_dot_general<1>;
-concat_pad<1>;
-reduce_pad<1>;
-broadcast_pad<1>;
-zero_product_reshape_pad<1>;
-mul_zero_pad<1>;
-div_zero_pad<1>;
-binop_const_reshape_pad<1>;
-binop_const_pad_add<1>;
-binop_const_pad_subtract<1>;
-binop_const_pad_mul<1>;
-binop_const_pad_div<1>;
-slice_reshape_pad<1>;
-binop_binop_pad_pad_add<1>;
-binop_binop_pad_pad_mul<1>;
-binop_pad_pad_add<1>;
-binop_pad_pad_subtract<1>;
-binop_pad_pad_mul<1>;
-binop_pad_pad_div<1>;
-binop_pad_pad_min<1>;
-binop_pad_pad_max<1>;
-unary_pad_push_convert<1>;
-unary_pad_push_tanh<1>;
-unary_pad_push_exp<1>;
-transpose_pad<1>;
-transpose_dot_reorder<1>;
-dot_transpose<1>;
-convert_convert_float<1>;
-concat_to_pad<1>;
-concat_appending_reshape<1>;
-reshape_iota<1>;
-broadcast_reduce<1>;
-slice_dot_general<1>;
-dot_reshape_pad<1>;
-pad_dot_general<1>(1);
-pad_dot_general<1>(0);
-},
-transform-interpreter,
-enzyme-hlo-remove-transform
-)""")
-
-def rmsnorm(x, weight):
-    ss = 1 / jnp.sqrt(x.dot(x) / x.shape[0] + 1e-5)
-    return weight * x * ss
-
-
-def softmax(x):
-    max_val = jnp.max(x)
-    x = jnp.exp(x - max_val)
-    return x / sum(x)
-
-
-def sigmoid(x):
-    return 1 / (1 + jnp.exp(-x))
-
-
-def silu(x):
-    return x * sigmoid(x)
-
-asserts = True
-
-def forward(x, config, weights, key_cache, value_cache):
-    pos = key_cache.shape[1]
-    assert pos == key_cache.shape[1]
-    assert pos == value_cache.shape[1]
-
-    n_layers = config["n_layers"]
-    seq_len = config["seq_len"]
-    n_heads = config["n_heads"]
-    vocab_size = config["vocab_size"]
-
-    # Total number of parameters of the recurrent state
-    dim = config["dim"]
-
-    n_kv_heads = config["n_kv_heads"]
-
-    # number of hidden dimensions?
-    hidden_dim = config["hidden_dim"]
-
-    # Number of parameters per head
-    head_size = dim // n_heads
-
-    # Number of heads per kv
-    kv_mul = n_heads // n_kv_heads
-
-    # Number of parameters in a kv
-    kv_dim = dim // n_heads * n_kv_heads
-
-    wo = weights["wo"]
-    if asserts:
-        if wo.shape != (n_layers, dim, dim):
-            print(
-                wo.shape,
-                weights,
-                (
-                    n_layers,
-                    dim,
-                    kv_dim,
-                    kv_mul,
-                    head_size,
-                    hidden_dim,
-                    n_kv_heads,
-                    vocab_size,
-                    n_heads,
-                    seq_len,
-                    n_layers,
-                ),
-            )
-        assert wo.shape == (n_layers, dim, dim)
-    rms_ffn_weight = weights["rms_ffn_weight"]
-    if asserts:
-        assert rms_ffn_weight.shape == (n_layers, dim)
-    w1 = weights["w1"]
-    if asserts:
-        assert w1.shape == (n_layers, hidden_dim, dim)
-    w3 = weights["w3"]
-    if asserts:
-        assert w3.shape == (n_layers, hidden_dim, dim)
-    w2 = weights["w2"]
-    if asserts:
-        assert w2.shape == (n_layers, dim, hidden_dim)
-
-    rms_att_weight = weights["rms_att_weight"]
-    if asserts:
-        assert rms_att_weight.shape == (n_layers, dim)
-
-    rms_final_weight = weights["rms_final_weight"]
-    if asserts:
-        assert rms_final_weight.shape == (dim,)
-    wcls = weights["wcls"]
-    if asserts:
-        assert wcls.shape == (vocab_size, dim)
-
-    # token_embedding_table = weights['token_embedding_table']
-    # if asserts: assert token_embedding_table.shape == (vocab_size, dim)
-
-    # x = token_embedding_table[token, :]
-    # if asserts: assert x.shape == (dim, )
-
-    wq = weights["wq"]
-    if asserts:
-        assert wq.shape == (n_layers, dim, dim)
-
-    wk = weights["wk"]
-    if asserts:
-        assert wk.shape == (n_layers, kv_dim, dim)
-
-    wv = weights["wv"]
-    if asserts:
-        assert wv.shape == (n_layers, kv_dim, dim)
-
-    toconv = []
-
-    for i in range(0, dim, 2):
-        freq = 1 / jnp.power(10000, (i % head_size) / head_size)
-        val = pos * freq
-        fcr = jnp.cos(val)
-        fci = jnp.sin(val)
-
-        rotM = jnp.array([[fcr, -fci], [fci, fcr]])
-        toconv.append(rotM)
-    toconv2 = toconv[: kv_dim // 2] + [jnp.eye(2)] * (dim // 2 - kv_dim // 2)
-
-    toconv = jnp.array(toconv)
-    toconv2 = jnp.array(toconv2)
-
-    keys2 = []
-    values2 = []
-    for l in range(n_layers):
-        xb = rmsnorm(x, rms_att_weight[l, :])
-        if asserts:
-            assert xb.shape == (dim,)
-
-        q = wq[l, :, :] @ xb
-        if asserts:
-            assert q.shape == (dim,)
-
-        k = wk[l, :, :] @ xb
-        if asserts:
-            assert q.shape == (kv_dim,)
-
-        v = wv[l, :, :] @ xb
-        if asserts:
-            assert q.shape == (kv_dim,)
-
-        q_tmp = jnp.reshape(q, (dim // 2, 2))
-        k_tmp = jnp.reshape(k, (dim // 2, 2))
-
-        # dim == head_size * n_heads
-
-        # Batched gemv
-        k = jnp.reshape(jnp.einsum("ijk,ik -> ij", toconv2, k_tmp), (dim,))
-        q = jnp.reshape(jnp.einsum("ijk,ik -> ij", toconv, q_tmp), (dim,))
-
-        key_cache_l = key_cache[l, :, :]
-        key_cache_l = jnp.append(key_cache_l, jnp.reshape(k, (1, dim)), axis=0)
-        value_cache_l = value_cache[l, :, :]
-        value_cache_l = jnp.append(value_cache_l, jnp.reshape(v, (1, dim)), axis=0)
-        keys2.append(key_cache_l)
-        values2.append(value_cache_l)
-
-        xbs2 = []
-        for h in range(n_heads):
-            q2 = q[head_size * h : head_size * (h + 1)]
-            if asserts:
-                assert q2.shape == (head_size,)
-
-            # For kv_mul consecutive heads, they share the same kv cache
-            # reshape key_cache last dim from (kv_dim,) to (kv_mul, head_size)
-            # generalized einsum reducing the last dim, the rest are batch
-            att = []
-
-            key_index = h // kv_mul
-
-            att = jnp.einsum(
-                "ij,j->i",
-                key_cache_l[:, key_index * head_size : (key_index + 1) * head_size],
-                q2,
-            )
-
-            att = att / jnp.sqrt(head_size)
-
-            att = softmax(att)
-
-            x_tmp = jnp.einsum(
-                "ij,i->j",
-                value_cache_l[:, key_index * head_size : (key_index + 1) * head_size],
-                att,
-            )
-
-            xbs2.append(x_tmp)
-
-        # Todo right concat
-        xb = jnp.concatenate(xbs2, axis=None)
-
-        xb2 = wo[l, :, :] @ xb
-        if asserts:
-            assert xb2.shape == (dim,)
-
-        x += xb2
-
-        # Rmsnorm and feedforward swiglu
-
-        xb = rmsnorm(x, rms_ffn_weight[l, :])
-
-        # Now for FFN in PyTorch we have: self.w2(F.silu(self.w1(x)) * self.w3(x))
-        # first calculate self.w1(x) and self.w3(x)
-
-        hb = w1[l, :, :] @ xb
-        hb2 = w3[l, :, :] @ xb
-
-        hb = silu(hb)
-
-        hb = hb * hb2
-
-        xb = w2[l, :, :] @ hb
-
-        x += xb
-
-    x = rmsnorm(x, rms_final_weight)
-    logits = wcls @ x
-
-    return x
-
-
-class Llama(absltest.TestCase):
-    def somewhat_close(self, a, b, eps):
-        c = jnp.abs(a - b)
-        return jnp.max(c) < eps
-
-    def test_llama_random(self):
-        config = {
-            "dim": 288,
-            "hidden_dim": 768,
-            "n_layers": 6,
-            "n_heads": 6,
-            "n_kv_heads": 6,
-            "vocab_size": 32000,
-            "seq_len": 256,
-        }
-
-        n_layers = config["n_layers"]
-        seq_len = config["seq_len"]
-        n_heads = config["n_heads"]
-        dim = config["dim"]
-        n_kv_heads = config["n_kv_heads"]
-        vocab_size = config["vocab_size"]
-        hidden_dim = config["hidden_dim"]
-        kv_dim = dim // n_heads * n_kv_heads
-        head_size = dim // n_heads
-
-        key = jax.random.PRNGKey(0)
-        weights = {}
-        dweights = {}
-
-        for name, shape in [
-            ("rms_att_weight", (n_layers, dim)),
-            ("wq", (n_layers, dim, n_heads * head_size)),
-            ("wk", (n_layers, dim, n_kv_heads * head_size)),
-            ("wv", (n_layers, dim, n_kv_heads * head_size)),
-            ("wo", (n_layers, dim, dim)),
-            ("rms_ffn_weight", (n_layers, dim)),
-            ("w1", (n_layers, hidden_dim, dim)),
-            ("w2", (n_layers, dim, hidden_dim)),
-            ("w3", (n_layers, hidden_dim, dim)),
-            ("rms_final_weight", (dim,)),
-            ("wcls", (vocab_size, dim)),
-        ]:
-            key, subkey = jax.random.split(key)
-            key, subkey2 = jax.random.split(key)
-            weights[name] = jax.random.uniform(subkey, shape=shape)
-            dweights[name] = jax.random.uniform(subkey2, shape=shape)
-
-        key, subkey = jax.random.split(key)
-        x = jax.random.uniform(subkey, shape=(dim,))
-        key, subkey = jax.random.split(key)
-        dx = jax.random.uniform(subkey, shape=(dim,))
-
-        def partial(func, config):
-            def sfn(x, weights, key_cache, value_cache):
-                return func(x, config, weights, key_cache, value_cache)
-
-            return sfn
-
-        pos = 1
-        key_cache = jnp.zeros((n_layers, pos, kv_dim))
-        value_cache = jnp.zeros((n_layers, pos, kv_dim))
-
-        key, subkey = jax.random.split(key)
-        dkc = jax.random.uniform(subkey, shape=(n_layers, pos + 1, kv_dim))
-        key, subkey = jax.random.split(key)
-        dvc = jax.random.uniform(subkey, shape=(n_layers, pos + 1, kv_dim))
-
-        func = partial(forward, config)
-
-        repeats = 1000
-
-        jax_func = jax.jit(func)
-        jax_res = jax_func(x, weights, key_cache, value_cache)
-
-        # enzyme_pipeline = enzyme_jax.NewXLAPipeline(mlirad=True)
-        # enzyme_pipeline = enzyme_jax.JaXPipeline()
-        # enzyme_pipeline = enzyme_jax.OldXLAPipeline()
-        # enzyme_pipeline = enzyme_jax.NewXLAPipeline(mlirad=False)
-        enzyme_pipeline = massive_pipeline
-
-        enzyme_func = jax.jit(
-            enzyme_jax.enzyme_jax_ir(argv=argv, pipeline_options=enzyme_pipeline)(func)
-        )
-        enzyme_res = enzyme_func(x, weights, key_cache, value_cache)
-        
-        eqsat_pipeline = enzyme_jax.JaXPipeline(
-            "inline{default-pipeline=canonicalize max-iterations=4},"
-            + "equality-saturation-pass"
-        )
-        eqsat_func = jax.jit(
-            enzyme_jax.enzyme_jax_ir(
-                argv=argv,
-                pipeline_options=eqsat_pipeline,
-            )(func)
-        )
-        eqsat_res = eqsat_func(x, weights, key_cache, value_cache)
-
-        eps = 1e-6
-        assert(self.somewhat_close(jax_res, enzyme_res, eps))
-        assert(self.somewhat_close(jax_res, eqsat_res, eps))
-
-        # print(
-        #     "Vanilla",
-        #     timeit.Timer(
-        #         "func(x, weights, key_cache, value_cache)",
-        #         globals={
-        #             "func": func,
-        #             "x": x,
-        #             "weights": weights,
-        #             "key_cache": key_cache,
-        #             "value_cache": value_cache,
-        #         },
-        #     ).timeit(repeats),
-        # )
-
-        print(
-            "jax primal",
-            timeit.Timer(
-                "jax_func(x, weights, key_cache, value_cache)",
-                globals={
-                    "jax_func": jax_func,
-                    "x": x,
-                    "weights": weights,
-                    "key_cache": key_cache,
-                    "value_cache": value_cache,
-                },
-            ).timeit(repeats),
-        )
-        
-        print(
-            "enzyme primal",
-            timeit.Timer(
-                "enzyme_func(x, weights, key_cache, value_cache)",
-                globals={
-                    "enzyme_func": enzyme_func,
-                    "x": x,
-                    "weights": weights,
-                    "key_cache": key_cache,
-                    "value_cache": value_cache,
-                },
-            ).timeit(repeats),
-        )
-        
-        print(
-            "eqsat primal",
-            timeit.Timer(
-                "eqsat_func(x, weights, key_cache, value_cache)",
-                globals={
-                    "eqsat_func": eqsat_func,
-                    "x": x,
-                    "weights": weights,
-                    "key_cache": key_cache,
-                    "value_cache": value_cache,
-                },
-            ).timeit(repeats),
-        ) 
-
-        if True:    
-            @jax.jit 
-            def jax_fwd(x, dx, weights, dweights, kc, dkc, vc, dvc):
-                return jax.jvp(jax_func, (x, weights, kc, vc), (x, weights, dkc, dvc))
-            jax_fwd_res = jax_fwd(x, dx, weights, dweights, key_cache, key_cache, value_cache, value_cache)
-
-            @jax.jit 
-            def enzyme_fwd(x, dx, weights, dweights, kc, dkc, vc, dvc):
-                return jax.jvp(enzyme_func, (x, weights, kc, vc), (x, weights, dkc, dvc))
-            enzyme_fwd_res = enzyme_fwd(x, dx, weights, dweights, key_cache, key_cache, value_cache, value_cache)
-
-            @jax.jit 
-            def eqsat_fwd(x, dx, weights, dweights, kc, dkc, vc, dvc):
-                return jax.jvp(eqsat_func, (x, weights, kc, vc), (x, weights, dkc, dvc))
-            eqsat_fwd_res = eqsat_fwd(x, dx, weights, dweights, key_cache, key_cache, value_cache, value_cache)
-
-            print(
-                "jax fwd",
-                timeit.Timer(
-                    "jax_fwd(x, dx, weights, dweights, key_cache, key_cache, value_cache, value_cache)",
-                    globals={
-                        "jax_fwd": jax_fwd,
-                        "x": x,
-                        "dx": dx,
-                        "weights": weights,
-                        "dweights": dweights,
-                        "key_cache": key_cache,
-                        "value_cache": value_cache,
-                    },
-                ).timeit(repeats),
-
-                "enzyme fwd",
-                timeit.Timer(
-                    "enzyme_fwd(x, dx, weights, dweights, key_cache, key_cache, value_cache, value_cache)",
-                    globals={
-                        "enzyme_fwd": enzyme_fwd,
-                        "x": x,
-                        "dx": dx,
-                        "weights": weights,
-                        "dweights": dweights,
-                        "key_cache": key_cache,
-                        "value_cache": value_cache,
-                    },
-                ).timeit(repeats),
-
-                "eqsat fwd",
-                timeit.Timer(
-                    "eqsat_fwd(x, dx, weights, dweights, key_cache, key_cache, value_cache, value_cache)",
-                    globals={
-                        "eqsat_fwd": eqsat_fwd,
-                        "x": x,
-                        "dx": dx,
-                        "weights": weights,
-                        "dweights": dweights,
-                        "key_cache": key_cache,
-                        "value_cache": value_cache,
-                    },
-                ).timeit(repeats),
-            )
-
-        if True:
-            @jax.jit
-            def jax_rev(x, weights, kc, vc, dx, dkc, dvc):
-                primals, f_vjp = jax.vjp(jax_func, x, weights, kc, vc)
-                return f_vjp(dx)
-            jax_rev_res = jax_rev(x, weights, key_cache, value_cache, dx, dkc, dvc)
-
-            jax_rev_with_enzyme = jax.jit(
-                enzyme_jax.enzyme_jax_ir(
-                    argv=argv,
-                    pipeline_options=enzyme_pipeline
-                )(jax_rev)
-            ) 
-            jax_rev_with_enzyme_res = jax_rev_with_enzyme(x, weights, key_cache, value_cache, dx, dkc, dvc)
-
-            jax_rev_with_eqsat = jax.jit(
-                enzyme_jax.enzyme_jax_ir(
-                    argv=argv,
-                    pipeline_options=eqsat_pipeline
-                )(jax_rev)
-            )
-            jax_rev_with_eqsat_res = jax_rev_with_eqsat(x, weights, key_cache, value_cache, dx, dkc, dvc)
-            
-            @jax.jit
-            def enzyme_rev(x, weights, kc, vc, dx, dkc, dvc):
-                primals, f_vjp = jax.vjp(enzyme_func, x, weights, kc, vc)
-                return f_vjp(dx)
-            enzyme_rev_res = enzyme_rev(x, weights, key_cache, value_cache, dx, dkc, dvc)
-            
-            @jax.jit
-            def eqsat_rev(x, weights, kc, vc, dx, dkc, dvc):
-                primals, f_vjp = jax.vjp(eqsat_func, x, weights, kc, vc)
-                return f_vjp(dx)
-            eqsat_rev_res = eqsat_rev(x, weights, key_cache, value_cache, dx, dkc, dvc)
-
-            print(
-                "jax rev",
-                timeit.Timer(
-                    "jax_rev(x, weights, key_cache, value_cache, dx, dkc, dvc)",
-                    globals={
-                        "jax_rev": jax_rev,
-                        "x": x,
-                        "weights": weights,
-                        "key_cache": key_cache,
-                        "value_cache": value_cache,
-                        "dx": dx,
-                        "dkc": dkc,
-                        "dvc": dvc,
-                    },
-                ).timeit(repeats),
-            )
-
-            print(
-                "jax rev with enzyme",
-                timeit.Timer(
-                    "jax_rev_with_enzyme(x, weights, key_cache, value_cache, dx, dkc, dvc)",
-                    globals={
-                        "jax_rev_with_enzyme": jax_rev_with_enzyme,
-                        "x": x,
-                        "weights": weights,
-                        "key_cache": key_cache,
-                        "value_cache": value_cache,
-                        "dx": dx,
-                        "dkc": dkc,
-                        "dvc": dvc,
-                    },
-                ).timeit(repeats),
-            )
-            
-            print(
-                "jax rev with eqsat",
-                timeit.Timer(
-                    "jax_rev_with_eqsat(x, weights, key_cache, value_cache, dx, dkc, dvc)",
-                    globals={
-                        "jax_rev_with_eqsat": jax_rev_with_eqsat,
-                        "x": x,
-                        "weights": weights,
-                        "key_cache": key_cache,
-                        "value_cache": value_cache,
-                        "dx": dx,
-                        "dkc": dkc,
-                        "dvc": dvc,
-                    },
-                ).timeit(repeats),
-            )
-
-            print(
-                "enzyme rev",
-                timeit.Timer(
-                    "enzyme_rev(x, weights, key_cache, value_cache, dx, dkc, dvc)",
-                    globals={
-                        "enzyme_rev": enzyme_rev,
-                        "x": x,
-                        "weights": weights,
-                        "key_cache": key_cache,
-                        "value_cache": value_cache,
-                        "dx": dx,
-                        "dkc": dkc,
-                        "dvc": dvc,
-                    },
-                ).timeit(repeats),
-            )
-
-            print(
-                "eqsat rev",
-                timeit.Timer(
-                    "eqsat_rev(x, weights, key_cache, value_cache, dx, dkc, dvc)",
-                    globals={
-                        "eqsat_rev": eqsat_rev,
-                        "x": x,
-                        "weights": weights,
-                        "key_cache": key_cache,
-                        "value_cache": value_cache,
-                        "dx": dx,
-                        "dkc": dkc,
-                        "dvc": dvc,
-                    },
-                ).timeit(repeats),
-            )
-
-if __name__ == "__main__":
-    absltest.main()
\ No newline at end of file
diff --git a/test/llama.py b/test/llama.py
index 9e6d4622..62e94465 100644
--- a/test/llama.py
+++ b/test/llama.py
@@ -315,7 +315,6 @@ def forward(x, config, weights, key_cache, value_cache):
 exp_simplify<16>;
 slice_simplify<16>;
 convert_simplify<16>;
-reshape_simplify<16>;
 dynamic_slice_to_static<16>;
 dynamic_update_slice_elim<16>;
 concat_to_broadcast<16>;
diff --git a/test/maxtext.py b/test/maxtext.py
new file mode 100644
index 00000000..a972a601
--- /dev/null
+++ b/test/maxtext.py
@@ -0,0 +1,193 @@
+# Steps for getting results here
+# Run:
+#   1) pip install https://github.com/wsmoses/maxtext
+#   2) bazel build -c opt //:wheel
+#   3) pip install ./bazel-bin/*whl
+#   4) python test/maxtext.py
+
+from absl.testing import absltest
+import jax.numpy as jnp
+import jax.random
+import jax.lax
+import enzyme_ad.jax as enzyme_jax
+from enzyme_ad.jax import (
+    enzyme_jax_ir,
+    NewXLAPipeline,
+    OldXLAPipeline,
+    JaXPipeline,
+    hlo_opts,
+)
+import numpy as np
+import timeit
+
+argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
+
+import jax.numpy as np
+import numpy as onp
+from jax import jit
+from jax import random
+from jax import lax
+
+partialopt = (
+    "inline{default-pipeline=canonicalize max-iterations=4},"
+    + """canonicalize,cse,
+enzyme-hlo-generate-td{
+            patterns=compare_op_canon<16>;
+transpose_transpose<16>;
+broadcast_in_dim_op_canon<16>;
+convert_op_canon<16>;
+dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
+chained_dynamic_broadcast_in_dim_canonicalization<16>;
+dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
+noop_reduce_op_canon<16>;
+empty_reduce_op_canon<16>;
+dynamic_reshape_op_canon<16>;
+get_tuple_element_op_canon<16>;
+real_op_canon<16>;
+imag_op_canon<16>;
+get_dimension_size_op_canon<16>;
+gather_op_canon<16>;
+reshape_op_canon<16>;
+merge_consecutive_reshapes<16>;
+transpose_is_reshape<16>;
+zero_extent_tensor_canon<16>;
+reorder_elementwise_and_shape_op<16>;
+
+cse_broadcast_in_dim<16>;
+cse_slice<16>;
+cse_transpose<16>;
+cse_convert<16>;
+cse_pad<16>;
+cse_dot_general<16>;
+cse_reshape<16>;
+cse_mul<16>;
+cse_div<16>;
+cse_add<16>;
+cse_subtract<16>;
+cse_min<16>;
+cse_max<16>;
+cse_neg<16>;
+cse_concatenate<16>;
+
+concatenate_op_canon<16>(1024);
+select_op_canon<16>(1024);
+add_simplify<16>;
+sub_simplify<16>;
+and_simplify<16>;
+max_simplify<16>;
+min_simplify<16>;
+or_simplify<16>;
+negate_simplify<16>;
+mul_simplify<16>;
+div_simplify<16>;
+rem_simplify<16>;
+pow_simplify<16>;
+sqrt_simplify<16>;
+cos_simplify<16>;
+sin_simplify<16>;
+noop_slice<16>;
+const_prop_through_barrier<16>;
+slice_slice<16>;
+shift_right_logical_simplify<16>;
+pad_simplify<16>;
+negative_pad_to_slice<16>;
+tanh_simplify<16>;
+exp_simplify<16>;
+slice_simplify<16>;
+convert_simplify<16>;
+dynamic_slice_to_static<16>;
+dynamic_update_slice_elim<16>;
+concat_to_broadcast<16>;
+reduce_to_reshape<16>;
+broadcast_to_reshape<16>;
+gather_simplify<16>;
+iota_simplify<16>(1024);
+broadcast_in_dim_simplify<16>(1024);
+convert_concat<1>;
+dynamic_update_to_concat<1>;
+slice_of_dynamic_update<1>;
+slice_elementwise<1>;
+slice_pad<1>;
+dot_reshape_dot<1>;
+concat_const_prop<1>;
+concat_fuse<1>;
+pad_reshape_pad<1>;
+pad_pad<1>;
+concat_push_binop_add<1>;
+concat_push_binop_mul<1>;
+scatter_to_dynamic_update_slice<1>;
+reduce_concat<1>;
+slice_concat<1>;
+
+bin_broadcast_splat_add<1>;
+bin_broadcast_splat_subtract<1>;
+bin_broadcast_splat_div<1>;
+bin_broadcast_splat_mul<1>;
+slice_reshape<1>;
+
+dot_reshape_pad<1>;
+pad_dot_general<1>(1);
+broadcast_reduce<1>;
+            },
+            transform-interpreter,
+            enzyme-hlo-remove-transform,cse"""
+)
+
+pipelines = [
+    ("JaX  ", None),
+    ("JaXPipe", JaXPipeline()),
+    (
+        "HLOOpt",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "canonicalize,cse,enzyme-hlo-opt,cse"
+        ),
+    ),
+    ("PartOpt", JaXPipeline(partialopt)),
+    ("DefOpt", JaXPipeline(hlo_opts())),
+    (
+        "EqSat",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "equality-saturation-pass"
+        ),
+    ),
+]
+
+
+class MaxText(absltest.TestCase):
+    def setUp(self):
+        import MaxText
+        import MaxText.pyconfig
+
+        MaxText.pyconfig.initialize(
+            [
+                None,
+                "test/maxtext_configs/base.yml",
+                "dataset_type=synthetic",
+                "steps=10",
+            ]
+        )
+
+    def test(self):
+        import MaxText
+        import MaxText.pyconfig
+        import MaxText.train
+
+        config = MaxText.pyconfig.config
+
+        for name, pipeline in pipelines:
+            print("name=", name)
+
+            def rewrite(fn):
+                if pipeline is None:
+                    return fn
+                else:
+                    return enzyme_jax_ir(pipeline_options=pipeline, argv=argv)(fn)
+
+            res1 = MaxText.train.train_loop(config, prejit=rewrite)
+            print("name=", name, res1)
+
+
+if __name__ == "__main__":
+    absltest.main()
diff --git a/test/maxtext_configs/base.yml b/test/maxtext_configs/base.yml
new file mode 100644
index 00000000..40bef0b5
--- /dev/null
+++ b/test/maxtext_configs/base.yml
@@ -0,0 +1,456 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This sentinel is a reminder to choose a real run name.
+# If there is already a checkpoint under this run, that checkpoint will auto-resume.
+run_name: ""
+
+model_name: "default" # override config settings to match a specific model. other than the override, nothing should use this!
+normalization_layer_epsilon: 1.e-05
+
+################################## CHECKPOINTING ##################################
+# Checkpointing makes the following choices in the following order, starting with (1):
+#   (1) If there is already a checkpoint for this run_name, we load the latest entire checkpoint.
+#     This ensures if we're resuming a run after preemption or hardware failure we lose minimum state.
+#   (2) Same priority and mutually exclusive -- you can't set both!
+#      * If load_parameters_path is set, we load a parameter only checkpoint from that path.
+#      * If load_full_state_path is set, we load a full state checkpoint from that path.
+#   (3) We don't load a checkpoint and initialize state instead!
+
+# Loads a just parameters from a specific directory
+# e.g. gs://my-base-output-directory/my-previous-run-name/checkpoints/items/NUMBER or NUMBER/items
+load_parameters_path: ""
+# Loads a full checkpoint including optimizer state and step count from a specific directory
+# e.g. gs://my-base-output-directory/my-previous-run-name/checkpoints/items/NUMBER or NUMBER/items
+load_full_state_path: ""
+
+# If enable_checkpointing is true, an asynchronous checkpointer will be used if
+# async_checkpointing is true, else a synchronous one is used. If you have
+# problems with the checkpointer we recommend trying the synchronous one.
+enable_checkpointing: True
+async_checkpointing: True
+checkpoint_period: 10_000
+# enables one replica to read the ckpt then broadcast to the rest
+enable_single_replica_ckpt_restoring: False
+
+force_unroll: False # during generate_param_only_checkpoint should we unroll the loop?
+
+# checkpointing using orbax has two important parameters: array driver
+# and its underlying storage - the kvstore (preferably ocdbt)
+# orbax supports setting a target file size, chunking a single 
+# large arrays into small physical files (<2GB) can speed up distributed and over 
+# the network loading enormously
+checkpoint_storage_target_data_file_size_bytes: 2147483648
+checkpoint_storage_use_ocdbt: True
+checkpoint_storage_use_zarr3: True
+############################### END CHECKPOINTING ##################################
+
+
+reuse_example_batch: 0 # for testing TPU performance, this options repeated uses the same batch.
+
+metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written.
+# If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
+gcs_metrics: False
+
+# If true save config to GCS in {base_output_directory}/{run_name}/
+save_config_to_gcs: False
+
+# Activation dtypes.
+dtype: "bfloat16"
+# Used to configure quantization in the transformer layers, defaults to null implying bf16.
+# Possible alternative settings are as follows:
+# 'int8' for dynamic range quantization using 8-bits
+# 'int8w' for weights only quantization using 8-bits
+# 'int4w' for weights only quantization using 4-bits
+# 'intmp' for mixed precision weight only quantization based on config file
+# 'fp8' for 8-bit floating-point GeMMs on NVIDIA GPUs.
+quantization: ""
+# Choose one of default, high, and highest.
+# https://kolonist26-jax-kr.readthedocs.io/en/latest/jax.lax.html#jax.lax.Precision
+matmul_precision: "default"
+activations_in_float32: False # Sets activations to float32 before nonlinearity it true, else dtype
+# Path to file with quantization config - only used for mixed precision.
+# Example configs in ../Maxtext/configs/quantization
+# Allows us to configure different bits, tiling and scale for quantizing selected weights.
+# Bits represents number of bits to quantize to,
+# tile-size represents the tiling sized used in AQT tiled_dot_general,
+# Value of scale is used to scale the abs_max value used for AQT quantization
+# Defaults values are 8 bits, tile-size=-1 (no tiling) and scale=1.
+quant_cfg_path: ""
+quantize_kvcache: False # Set to True to quantize KV Cache values, defaults to False
+# Valid kv_quant_axis values:
+#   - "" is valid only when quantize_kvcache is False
+#   - "dkv" indicates quantize kv cache over the cache_kv, i.e. kv dimension axis
+#   - "heads_and_dkv" indicates quantize kv cache over cache_heads and cache_kv axes
+# Default to "heads_and_dkv" for faster compution, kv_quant_axis is not used when quantize_kvcache is False
+#   - "dkv" is expected with better accuracy but degraded computation
+kv_quant_axis: "heads_and_dkv"
+kv_quant_dtype: "int8"
+checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint
+# Saves params quantized on fly at following path
+save_quantized_params_path: ""
+
+# Shard the range finding operation for quantization. By default this is set to number of slices.
+quantization_local_shard_count: -1
+
+decoder_block: "llama2" # which style of DecoderBlock to use.
+# Global parameter scale needs to be a power of 2. If you want finer grained control of the model sizes
+# then you should explicitly set base_embed_dim, base_num_query_heads, base_num_kv_heads,
+# base_mlp_dim, base_num_decoder_layers and/or head_dim.
+weight_dtype: float32
+global_parameter_scale: 1
+base_emb_dim: 512
+# base_emb_dim: 2048
+base_num_query_heads: 1
+# base_num_query_heads: 16
+base_num_kv_heads: 1
+# base_num_kv_heads: 16
+base_mlp_dim: 7168
+base_num_decoder_layers: 1
+# base_num_decoder_layers: 16
+head_dim: 128
+mlp_activations: ["silu", "linear"]
+dropout_rate: 0
+logits_via_embedding: False
+normalize_embedding_logits: True  # whether to normlize pre-softmax logits if logits_via_embedding is true
+logits_dot_in_fp32: False  # whether to use fp32 in logits_dense or shared_embedding dot product for stability
+cast_logits_to_fp32: True # whether to cast the logits to fp32. The higher precision is generally beneficial, but it can vary slightly.
+
+# mixture of experts (moe)
+num_experts: 1
+num_experts_per_tok: 1
+megablox: True
+capacity_factor: -1.0 # a factor to decide expert capacity for token dropping, and no dropping by default
+load_balance_loss_weight: 0.01 # weight for the load balance loss
+
+# pipeline parallelism
+# The number of decoder layers is equal to the product of num_stages, num_layers_per_pipeline_stage and num_pipeline_repeats.
+# There is a tradeoff between the num_layers_per_pipeline_stage and num_pipeline_repeats: The more layers per stage the easier
+# it is to hide the pipeline communication behind the compute since there is more compute per stage, however there will be a larger bubble
+# since there are fewer repeats. Similarly there is tradeoff for num_pipeline_microbatches - more microbatches leads to a smaller bubble,
+# but a smaller size per microbatch which may hurt per-stage performance. Additionally note when microbatches > num_stages we have the opportunity to
+# perform the circular transfer (last stage to first) asynchronously.
+# The bubble fraction is (num_stages - 1) / (num_pipeline_repeats * num_pipeline_microbatches + num_stages - 1)
+num_layers_per_pipeline_stage: 1
+# The number of repeats will be set to num_decoder_layers / (num_pipeline_stages * num_layers_per_pipeline_stage)
+num_pipeline_repeats: -1
+# num_pipeline_microbatches must be a multiple of the number of pipeline stages. By default it is set to the number of stages.
+# Note the microbatch_size is given by global_batch_size / num_pipeline_microbatches, where global_batch_size = per_device_batch_size * num_devices
+num_pipeline_microbatches: -1
+scan_pipeline_iterations: True # This can be set independently of scan_layers, which is relevant when num_layers_per_pipeline_stage > 1.
+pipeline_delay_activation_forwarding: False # This delays the activation forwarding one loop iteration simplifying XLA's task of overlapping since
+# the communication and compute in each iteration are now independent. However this comes at the cost of doubling the pipeline bubble,
+# and you must set the number of microbatches to at least 2 * num_stages (the minimum 2 * num_stages is set by default with this delay).
+
+# Choose 'remat_policy' between 'minimal', 'save_dot_except_mlpwi', 'save_dot_except_mlp', 'save_qkv_proj', 'qkv_proj_offloaded', 'minimal_offloaded', 'save_out_proj' and 'full'.
+# These options offer a trade-off between speed (fastest to slowest) and HBM usage (highest to lowest)
+remat_policy: 'full'
+scan_layers: True
+param_scan_axis: 1
+
+# The attention parameter dictates the specific algorithm/methodology used to compute the attention scores
+# The attention_type parameter determines the variants of attention, e.g. global or local_sliding
+attention: 'dot_product' # Supported attention: autoselected, dot_product, flash, cudnn_flash_te
+attention_type: 'global' # Supported attention_type: global, local_sliding
+sliding_window_size: 0
+attn_logits_soft_cap: 0.0
+final_logits_soft_cap: 0.0
+use_post_attn_norm: False
+use_post_ffw_norm: False
+
+
+# Combine matmuls for QKV and MLP
+fused_qkv: False
+fused_mlp: False
+
+record_internal_nn_metrics: 0
+
+# Output directory
+# Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/"
+base_output_directory: ""
+
+# Whether or not to enable emergency checkpoint. If True, `local_checkpoint_directory` and a non-zero `local_checkpoint_period` must also be specified.
+# Emergency checkpoint is an experimental Orbax feature that: periodically saves to persistent storage and, with a larger invertal, saves to a local directory.
+# During restore, if a local copy is available in any slice, it will be broadcast to other slices without having to fetch from persistent storage.
+# See more details on https://github.com/google/orbax/tree/main/checkpoint/orbax/checkpoint/experimental/emergency.
+enable_emergency_checkpoint: False
+
+# It should be specified when and only when `enable_emergency_checkpoint` is True.
+local_checkpoint_directory: ""
+
+# It should be a positive number when and only when `enable_emergency_checkpoint` is True.
+local_checkpoint_period: 0
+
+# Jax cache directory
+jax_cache_dir: "~/jax_cache"
+
+# Hardware
+hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu', 'gpu_multiprocess' and 'cpu'
+
+# Parallelism
+mesh_axes: ['data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'tensor', 'expert', 'autoregressive']
+logical_axis_rules: [
+                      ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
+                      ['activation_batch_no_exp', ['data', 'fsdp', 'fsdp_transpose']],
+                       # For pipeline parallelism the pre and post decoder layer tensors' batch dimension is sharded by stages.
+                       # Microbatches are sharded by stage, so moving out of and into this sharding should be a local reshape.
+                       # The "stage" needs to be listed first since the microbatch dimension is first before the reshape.
+                      ['activation_embed_and_logits_batch', ['stage', 'data', 'fsdp', 'fsdp_transpose', 'expert']],
+                      ['activation_heads', ['tensor','sequence']],
+                      ['activation_kv_heads', ['tensor','sequence']],
+                      ['activation_length', 'sequence'],
+                      ['activation_embed', 'tensor'],
+                      ['activation_mlp', 'tensor'],
+                      ['activation_kv', 'tensor'],
+                      ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
+                      ['activation_kv_head_dim', 'tensor'],
+                      ['activation_vocab', ['tensor', 'sequence']],
+                      ['activation_vocab', 'tensor'],
+                      ['activation_vocab', 'sequence'],
+                      ['activation_stage', 'stage'],
+                      ['activation_exp', 'expert'],
+                      ['mlp', ['fsdp_transpose', 'tensor', 'autoregressive']],
+                      ['vocab', ['tensor', 'autoregressive']],
+                      ['embed', ['fsdp', 'fsdp_transpose', 'sequence', 'expert']],
+                      ['embed', ['fsdp', 'sequence', 'expert']],
+                      ['embed_no_exp', ['fsdp', 'fsdp_transpose', 'sequence']],
+                      ['embed_no_exp', ['fsdp', 'sequence']],
+                      ['norm', 'tensor'],
+                      ['heads', ['tensor', 'autoregressive']],
+                      ['layers', 'stage'],
+                      ['kv', []],
+                      ['kv_heads', ['tensor', 'autoregressive']],
+                      ['kv_head_dim', []],
+                      ['cache_batch', []],
+                      ['cache_heads', ['autoregressive', 'tensor']],
+                      ['cache_kv', []],
+                      ['cache_sequence', []],
+                      ['exp', 'expert'],
+                    ]
+# Axes used for DCN must be earlier in this list than ICI, see (b/339009148) for details
+data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'tensor', 'expert', 'autoregressive']]
+
+# One axis for each parallelism type may hold a placeholder (-1)
+# value to auto-shard based on available slices and devices.
+# By default, product of the DCN axes should equal number of slices
+# and product of the ICI axes should equal number of devices per slice.
+dcn_data_parallelism: -1  # recommended DCN axis to be auto-sharded
+dcn_fsdp_parallelism: 1
+dcn_fsdp_transpose_parallelism: 1
+dcn_sequence_parallelism: 1  # never recommended
+dcn_tensor_parallelism: 1 # never recommended
+dcn_pipeline_parallelism: 1
+dcn_expert_parallelism: 1
+dcn_autoregressive_parallelism: 1 # never recommended
+ici_data_parallelism: 1
+ici_fsdp_parallelism: -1 # recommended ICI axis to be auto-sharded
+ici_fsdp_transpose_parallelism: 1
+ici_sequence_parallelism: 1
+ici_tensor_parallelism: 1
+ici_autoregressive_parallelism: 1
+ici_pipeline_parallelism: 1
+ici_expert_parallelism: 1
+
+# The number of TPU slices is automatically determined, you should not set this explicitly. For ahead of time compilation,
+# you should set compile_toplogy_num_slices, which will in turn set this value. For non-TPU environments this is set to 1.
+num_slices: -1
+
+# Tokenizer
+vocab_size: 32_000 # powers of 2 for sharding
+tokenizer_path: "assets/tokenizer.llama2"
+tokenize_train_data: True  # False if the dataset is pre-tokenized
+tokenize_eval_data: True  # False if the dataset is pre-tokenized
+add_bos: True
+add_eos: True
+
+# Dataset
+per_device_batch_size: 12.0
+expansion_factor_real_data: -1 # if -1 then all hosts will load real data, else total_hosts//expansion_factor_real_data will pull data from GCS.
+eval_per_device_batch_size: 0
+max_corpus_chars: 10_000_000
+train_data_column: 'text'
+eval_data_column: 'text'
+# dataset_type must be synthetic, hf, grain, tfds
+# details in: https://github.com/google/maxtext/blob/main/getting_started/Data_Input_Pipeline.md
+dataset_type: tfds
+# for TFDS input pipeline (dataset_type=tfds)
+dataset_path: "" # your path given as argument in download_dataset.sh, e.g. "gs://my-maxtext-dataset/"
+dataset_name: 'c4/en:3.0.1'
+eval_dataset_name: 'c4/en:3.0.1'
+eval_split: 'validation'
+# for HuggingFace input pipeline (dataset_type=hf)
+hf_path: ''
+hf_data_dir: ''
+hf_train_files: ''
+hf_eval_split: ''
+hf_eval_files: ''
+hf_access_token: ''
+# for Grain input pipeline (dataset_type=grain)
+grain_train_files: ''
+grain_eval_files: ''
+grain_worker_count: 1
+
+# Training loop
+steps: 150_001 # If set to -1 then will inherit value from learning_rate_schedule_steps
+log_period: 100 # Flushes Tensorboard
+
+# We take inspiration from Llama2's learning rate (LR) schedule, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
+# Learning rate schedule has either two or three parts:
+# 1) Linear warmup from 0 to [learning_rate] over steps 0 to [learning_rate_schedule_steps * warmup_steps_fraction]
+# 2) Cosine decay from [learning_rate] to [learning_rate * cosine_learning_rate_final_fraction] from warmup to learning_rate_schedule_steps
+# 3) Constant learning rate of 0 from learning_rate_schedule_steps to steps.
+# The zero learning rate section can be used to more accurately measure the fully trained model's performance.
+learning_rate: 3.e-5
+cosine_learning_rate_final_fraction: 0.1
+warmup_steps_fraction: 0.1
+learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.
+# However you may choose a longer schedule (learning_rate_schedule_steps > steps), in which case the training will end before
+# dropping fully down. Or you may choose a shorter schedule, where the unspecified steps will have a learning rate of 0.
+
+max_target_length: 2048 # Maximum sequence length
+max_prefill_predict_length: 64 # Maximum length for the prefill when doing autoregression
+prompt: "I love to" # Prompt for language model sampling.
+load_from_prefill_dir: False # If true, decode.py doesn't "prefill" but just reads from directory
+prefill_cache_dir: "" # If set and load_from_prefill_dir, decode.py reads from directory. If set, decode.py writes to directory
+autoregressive_decode_assert: ""
+
+# For nsys profiler, pass the training command to nsys command
+# e.g. nsys profile -s none --force-overwrite true --capture-range=cudaProfilerApi --capture-range-end=stop {training command}
+profiler: "" # Supported profiler: '', xplane, nsys
+# If set to true, upload all profiler results from all hosts. Otherwise, only upload the profiler result from the first host.
+upload_all_profiler_results: False
+# Skip first n steps for profiling, to omit things like compilation and to give
+# the iteration time a chance to stabilize.
+skip_first_n_steps_for_profiler: 1
+# Profile for a small number of steps to avoid a large profile file size.
+profiler_steps: 5
+
+# When dropout is false the model is a deterministic function of the
+# data_shuffle_seed and init_weights_seed (i.e. reproducible losses)
+enable_dropout: True
+enable_data_shuffling: True
+data_shuffle_seed: 0
+init_weights_seed: 0
+
+# You may disable clipping by setting gradient_clipping_threshold to zero.
+gradient_clipping_threshold: 1.0
+
+# Instead of updating the weights every step, you may effectively use a larger
+# batch by accumulating the gradient over a set of steps.
+gradient_accumulation_steps: 1
+
+# AdamW optimizer parameters
+# We use AdamW following Llama2's training details, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
+opt_type: "adamw"  # one of "adam_pax" or "adamw"
+adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients.
+adam_b2: 0.95 # Exponential decay rate to track the second moment of past gradients.
+adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root.
+adam_eps_root: 0. # A small constant applied to denominator inside the square root.
+adam_weight_decay: 0.1 # AdamW Weight decay
+
+# Stack trace parameters
+collect_stack_trace: False
+stack_trace_to_cloud: False  # Uploads to cloud logging if True, else to the console if False.
+stack_trace_interval_seconds: 600  # Stack trace collection frequency in seconds.
+
+# Use iota operator in Embed
+use_iota_embed: False
+# use positional embedding
+use_untrainable_positional_embedding: False
+trainable_position_size: -1  # enable gpt3 position embedding with a positive trainable_position_size
+# Rope parameters
+rope_min_timescale: 1
+rope_max_timescale: 10_000
+
+# Ahead of time Compilation (aka AOT)
+# Only set these arguments if you are running train_compile or loading a compiled train step.
+compiled_trainstep_file: "" # Name of saved serialized compiled train_step, e.g. compiled_train_v5e-256.pickle
+compile_topology: '' # Target hardware version, e.g. 'v5e-256'
+compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
+
+decode_sampling_strategy: "greedy" # decode_sampling_strategy should be one of greedy, weighted, nucleus, or topk
+decode_sampling_nucleus_p: -1 # set if you're doing nucleus / top-p
+decode_sampling_top_k: 0 # set if you're doing top-k
+decode_sampling_temperature: 1.
+
+eval_interval: -1  # the specific number of train step between eval_step
+eval_steps: -1  # only run this number of batches for eval, for debugging use
+target_eval_loss: 0.  # early stop once reaching target eval_loss
+
+# Goodput parameters
+enable_goodput_recording: False
+monitor_goodput: False
+goodput_upload_interval_seconds: 60
+
+# Vertex AI Tensorboard Configurations - https://github.com/google/maxtext/tree/main/getting_started/Use_Vertex_AI_Tensorboard.md
+# Set to True for GCE, False if running via XPK
+use_vertex_tensorboard: False
+# Project to create Vertex AI Tensorboard in for GCE, blank if project is set using 'gcloud config set project'
+# Set this to blank if running via XPK
+vertex_tensorboard_project: ""
+# Region to create Vertex AI Tensorboard in for GCE, blank if running via XPK
+# Vertex AI supported regions: https://cloud.google.com/vertex-ai/docs/general/locations#available-regions
+vertex_tensorboard_region: ""
+
+# If set to True, MaxText will perform extra checks using jax.checkify. Note that this will effect performance.
+max_checkify: False
+
+# Inference
+inference_microbenchmark_prefill_lengths: "64,128,256,512,1024"
+inference_microbenchmark_stages: "prefill,generate"
+inference_microbenchmark_loop_iters: 10
+inference_microbenchmark_log_file_path: ""
+inference_metadata_file: "" # path to a json file
+enable_model_warmup: False
+
+
+# KV Cache layout control
+# Logical layout: 0,1,2,3 ; CACHE_BATCH, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV
+# Default layout: 1,2,0,3 ; CACHE_SEQUENCE, CACHE_HEADS, CACHE_BATCH, CACHE_KV
+prefill_cache_axis_order: "1,2,0,3"
+ar_cache_axis_order: "1,2,0,3"
+
+# Compute layout control
+# Default layout: 0,1,2,3 ; BATCH, LENGTH, HEAD, D_KV
+# Currently only support compute layout: 0,1,2,3 and 0,2,1,3
+compute_axis_order: "0,1,2,3"
+
+reshape_q: False
+
+# Maxengine Metrics
+prometheus_port: 0
+
+# Maxengine server
+enable_jax_profiler: False
+jax_profiler_port: 9999
+
+# Checkpoint Structured logging
+enable_checkpoint_cloud_logger: False
+enable_checkpoint_standard_logger: False
+
+# Single-controller
+enable_single_controller: True
+
+# Split physical axes for https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.mesh_utils.create_device_mesh.html
+allow_split_physical_axes: False
+
+use_ragged_attention: False
+ragged_block_size: 256
+
+### Splash attention block sizes
+# These can be tuned for specific hardware generations, and can be set up to
+# the model's sequence length.
+sa_block_q: 512
+sa_block_q_dkv: 512
+sa_block_q_dq: 512
diff --git a/test/mistral.py b/test/mistral.py
new file mode 100644
index 00000000..16a03187
--- /dev/null
+++ b/test/mistral.py
@@ -0,0 +1,102 @@
+from transformers import FlaxMistralModel, MistralConfig, BertTokenizer
+import jax
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import absltest
+import jax.random
+from enzyme_ad.jax import JaXPipeline, hlo_opts
+from test_utils import *
+import llama
+from typing import Tuple
+import jax.numpy as jnp
+
+def create_custom_mistral_config(
+    vocab_size=50257,
+    hidden_size=768,
+    num_hidden_layers=12,
+    num_attention_heads=12,
+    intermediate_size=3072,
+    max_position_embeddings=2048,
+    hidden_dropout_prob=0.1,
+    attention_probs_dropout_prob=0.1,
+    initializer_range=0.02,
+    **kwargs
+):
+    return MistralConfig(
+        vocab_size=vocab_size,
+        hidden_size=hidden_size,
+        num_hidden_layers=num_hidden_layers,
+        num_attention_heads=num_attention_heads,
+        intermediate_size=intermediate_size,
+        max_position_embeddings=max_position_embeddings,
+        hidden_dropout_prob=hidden_dropout_prob,
+        attention_probs_dropout_prob=attention_probs_dropout_prob,
+        initializer_range=initializer_range,
+        **kwargs
+    )
+
+def create_custom_mistral_model(config: MistralConfig, input_shape: Tuple = (1, 1), seed: int = 0):
+    return FlaxMistralModel(config, input_shape=input_shape, seed=seed)
+
+pipelines = [
+    ("JaX", None, CurBackends),
+    ("JaXPipe", JaXPipeline(), CurBackends),
+    (
+        "HLOOpt",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "canonicalize,cse,enzyme-hlo-opt,cse"
+        ),
+        CurBackends,
+    ),
+    ("PartOpt", JaXPipeline(llama.partialopt), CurBackends),
+    ("DefOpt", JaXPipeline(hlo_opts()), CurBackends),
+    (
+        "EqSat",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "equality-saturation-pass"
+        ),
+        CurBackends,
+    ),
+]
+
+mistral_config = create_custom_mistral_config(
+    hidden_size=512,
+    num_hidden_layers=8,
+    num_attention_heads=8,
+)
+
+model = create_custom_mistral_model(mistral_config)
+
+# Load the tokenizer (maybe we should try a different one)
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+
+def mistral_forward(input_ids, attention_mask):
+    """ Mistral forward pass """
+    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+    return outputs.last_hidden_state
+
+class MistralTransformerTest(EnzymeJaxTest):
+    def setUp(self):
+        text = "This is a sample text for Mistral transformer test."
+        inputs = tokenizer(text, return_tensors="np")
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+
+        self.hidden_state = mistral_forward(input_ids, attention_mask)
+
+        self.fn = mistral_forward
+        self.name = "mistral"
+        self.count = 1000
+        self.revprimal = False
+        self.AllPipelines = pipelines
+        self.AllBackends = CurBackends
+
+        self.ins = [input_ids, attention_mask]
+        self.dins = [input_ids, attention_mask]
+        self.douts = [self.hidden_state]
+        self.tol = 5e-5
+
+if __name__ == "__main__":
+    absltest.main()
diff --git a/test/resnet.py b/test/resnet.py
new file mode 100644
index 00000000..c45fc365
--- /dev/null
+++ b/test/resnet.py
@@ -0,0 +1,86 @@
+from transformers import AutoImageProcessor, FlaxResNetModel, ResNetConfig
+from PIL import Image
+import requests
+import jax
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import absltest
+import jax.random
+from enzyme_ad.jax import JaXPipeline, hlo_opts, enzyme_jax_ir
+from test_utils import *
+import llama
+
+# Define the pipelines as per your framework
+pipelines = [
+    ("JaX", None, CurBackends),
+    ("JaXPipe", JaXPipeline(), CurBackends),
+    (
+        "HLOOpt",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "canonicalize,cse,enzyme-hlo-opt,cse"
+        ),
+        CurBackends,
+    ),
+    ("PartOpt", JaXPipeline(llama.partialopt), CurBackends),
+    ("DefOpt", JaXPipeline(hlo_opts()), CurBackends),
+    (
+        "EqSat",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "equality-saturation-pass"
+        ),
+        CurBackends,
+    ),
+]
+
+# Load the image processor and model
+processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
+model = FlaxResNetModel.from_pretrained("microsoft/resnet-50")
+
+# Define the ResNet forward pass
+def resnet_forward(pixel_values, weights):
+    """ ResNet forward pass """
+    model.params = weights  # Set the weights to the loaded weights
+    outputs = model(pixel_values=pixel_values)
+    return outputs.last_hidden_state
+
+# Define the test class for ResNet using absltest
+class ResNetTest(EnzymeJaxTest):
+    def setUp(self):
+        # Create a ResNetConfig object
+        self.config = ResNetConfig.from_pretrained("microsoft/resnet-50")
+
+        # Load the weights for ResNet model (Microsoft pretrained)
+        self.weights = model.params
+
+        # Sample image from COCO dataset
+        url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        pil_image = Image.open(requests.get(url, stream=True).raw)
+        np_image = np.array(pil_image)
+
+        # Process the image
+        inputs = processor(images=np_image, return_tensors="np")
+        pixel_values = inputs["pixel_values"]  # Extract the actual NumPy array
+
+        # Run the ResNet forward pass with the processed input
+        self.logits = resnet_forward(pixel_values, self.weights)
+
+        # Setup for the test harness
+        self.fn = resnet_forward
+        self.name = "resnet"
+        self.count = 1000
+        self.revprimal = False
+        self.AllPipelines = pipelines
+        self.AllBackends = CurBackends
+
+        # Input and output setup
+        self.ins = [pixel_values, self.weights]
+        self.dins = [pixel_values, self.weights]
+        self.douts = [self.logits]
+        self.tol = 5e-5
+
+
+# Running the tests with absltest
+if __name__ == "__main__":
+    absltest.main()
diff --git a/test/simpler.py b/test/simpler.py
index d893d33c..ca132e77 100644
--- a/test/simpler.py
+++ b/test/simpler.py
@@ -15,9 +15,9 @@ def test_simple_random(self):
         efunc = enzyme_jax.enzyme_jax_ir(pipeline_options=enzyme_jax.JaXPipeline("equality-saturation-pass"),)(test)
         
         ka, kb, kc, kd = jax.random.split(jax.random.PRNGKey(0), num=4)
-        a = jax.random.uniform(ka, shape=(10, 10))
-        b = jax.random.uniform(kb, shape=(10, 10))
-        c = jax.random.uniform(kc, shape=(10, 10))
+        a = jax.random.uniform(ka, shape=(3, 10, 10))
+        b = jax.random.uniform(kb, shape=(3, 10, 10))
+        c = jax.random.uniform(kc, shape=(3, 10, 10))
         d = jax.random.uniform(kd, shape=(20, 500))
         e = jax.random.uniform(ka, shape=(20, 500))
         f = jax.random.uniform(ka, shape=(5, 5))
diff --git a/test/test_utils.py b/test/test_utils.py
index d340d76b..0d1b6021 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -35,10 +35,22 @@ def timeit(filename, pipeline, stage, str, globals, count, warmup=100):
 argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
 
 devices = []
-CurBackends = [jax.default_backend()]
+CurBackends = []
+backends_initialized = False
+
+
+def setup_backends():
+    global backends_initialized
+    global devices
+    global CurBackends
+    if backends_initialized:
+        return
+    backend = jax.default_backend()
+    CurBackends.append(backend)
+    if jax.default_backend() != "cpu":
+        devices.append(backend)
+    backends_initialized = True
 
-if jax.default_backend() != "cpu":
-    devices = CurBackends
 
 AllBackends = ["cpu"] + devices
 AllPipelines = [
@@ -154,6 +166,8 @@ def __init__(self, *args, **kwargs):
         self.AllPipelines = AllPipelines
         self.revprimal = True
         self.tol = 1e-6
+        self.mlirad_fwd = True
+        self.mlirad_rev = True
 
     def setUp(self):
         self.name = None
@@ -161,13 +175,12 @@ def setUp(self):
     def test(self):
         if self.name is None:
             return
+        setup_backends()
         self.harness(self.name, self.fn, self.ins, self.dins, self.douts)
 
     def harness(self, name, in_fn, ins, dins, douts):
         assert len(ins) == len(dins)
 
-        assert 1 == len(douts)
-
         primalstr = "fn(" + (", ".join(["in" + str(i) for i in range(len(ins))])) + ")"
 
         fwdstr = (
@@ -193,8 +206,10 @@ def harness(self, name, in_fn, ins, dins, douts):
             fwdins = primalins | {
                 ("din" + str(i)): dins_backend[i] for i in range(len(dins))
             }
-            revins = primalins | {"dout": douts_backend[0]}
-
+            if len(douts) == 1:
+                revins = primalins | {"dout": douts_backend[0]}
+            else:
+                revins = primalins | {"dout": tuple(douts_backend)}
             primres = None
 
             for pname, pipeline, pbackends in self.primfilter(self.AllPipelines):
@@ -233,40 +248,40 @@ def harness(self, name, in_fn, ins, dins, douts):
 
             for pname, pipeline, pbackends in self.fwdfilter(self.AllPipelines):
                 if backend in pbackends:
-                    rfn_enzyme = (
-                        in_fn
-                        if pipeline is None
-                        else jax.jit(
-                            enzyme_jax_ir(pipeline_options=pipeline, argv=argv)(in_fn),
+                    if self.mlirad_fwd or pipeline is None:
+                        rfn_enzyme = in_fn
+                        fwd_enzyme = jax.jit(
+                            (
+                                splatjvp(rfn_enzyme)
+                            if pipeline is None
+                            else enzyme_jax_ir(
+                                    pipeline_options=pipeline, argv=argv
+                            )(splatjvp(rfn_enzyme))
+                            ),
                             # backend=backend
                         )
-                    )
-                    fwd_enzyme = jax.jit(
-                        splatjvp(rfn_enzyme),
-                        # backend=backend
-                    )
 
-                    primals, tangents = fwd_enzyme(*(ins_backend + dins_backend))
+                        primals, tangents = fwd_enzyme(*(ins_backend+dins_backend))
 
-                    recursive_check(self, primals, primres, self.tol)
+                        recursive_check(self, primals, primres, self.tol)
 
-                    if fwdres is None:
-                        fwdres = tangents
-                    else:
-                        recursive_check(self, tangents, fwdres, self.tol)
+                        if fwdres is None:
+                            fwdres = tangents
+                        else:
+                            recursive_check(self, tangents, fwdres, self.tol)
 
-                    print(
-                        name,
-                        ",",
-                        pname,
-                        ",",
-                        backend,
-                        ",",
-                        "Forward",
-                        ",",
-                        timeit(csv_filename, pname, "Forward", fwdstr, {'fwd': fwd_enzyme} | fwdins, self.count),
-                        sep="\t",
-                    )
+                        print(
+                            name,
+                            ",",
+                            pname,
+                            ",",
+                            backend,
+                            ",",
+                            "Forward",
+                            ",",
+                            timeit(csv_filename, pname, "Forward", fwdstr, {'fwd': fwd_enzyme} | fwdins, self.count),
+                            sep="\t",
+                        )
 
             # assert fwdres is not None
 
@@ -276,45 +291,50 @@ def harness(self, name, in_fn, ins, dins, douts):
 
             for pname, pipeline, pbackends in self.revfilter(self.AllPipelines):
                 if backend in pbackends:
+
+                    adout = douts_backend
+                    if len(douts) != 1:
+                        adout = [tuple(douts_backend)]
                     if pipeline is not None:
-                        rfn_enzyme = (
-                            in_fn
-                            if pipeline is None
-                            else enzyme_jax_ir(pipeline_options=pipeline, argv=argv)(
+                        if self.mlirad_rev or pipeline is None:
+                            rfn_enzyme = (
                                 in_fn
+                                if pipeline is None
+                                else enzyme_jax_ir(
+                                    pipeline_options=pipeline, argv=argv
+                                )(in_fn)
+                            )
+                            rev_enzyme = jax.jit(
+                                revtransform(rfn_enzyme),
+                                # backend=backend
                             )
-                        )
-                        rev_enzyme = jax.jit(
-                            revtransform(rfn_enzyme),
-                            # backend=backend
-                        )
-
-                        if self.revprimal:
-                            primals, grads = rev_enzyme(*douts_backend, *ins_backend)
-                        else:
-                            grads = rev_enzyme(*douts_backend, *ins_backend)
-                            assert grads is not None
-
-                        if self.revprimal and primres is not None:
-                            recursive_check(self, primals, primres, self.tol)
-
-                        if revres is None:
-                            revres = grads
-                        else:
-                            recursive_check(self, grads, revres, self.tol)
 
-                        print(
-                            name,
-                            ",",
-                            pname,
-                            ",",
-                            backend,
-                            ",",
-                            "PreRev",
-                            ",",
-                            timeit(csv_filename, pname, "PreRev", revstr, {'rev': rev_enzyme} | revins, self.count),
-                            sep="\t",
-                        )
+                            if self.revprimal:
+                                primals, grads = rev_enzyme(*adout, *ins_backend)
+                            else:
+                                grads = rev_enzyme(*adout, *ins_backend)
+                                assert grads is not None
+
+                            if self.revprimal and primres is not None:
+                                recursive_check(self, primals, primres, self.tol)
+
+                            if revres is None:
+                                revres = grads
+                            else:
+                                recursive_check(self, grads, revres, self.tol)
+
+                            print(
+                                name,
+                                ",",
+                                pname,
+                                ",",
+                                backend,
+                                ",",
+                                "PreRev",
+                                ",",
+                                timeit(csv_filename, pname, "PreRev", revstr, {'rev': rev_enzyme} | revins, self.count),
+                                sep="\t",
+                            )
 
                         rfn_enzyme = in_fn
                         rev_enzyme = jax.jit(
@@ -329,9 +349,9 @@ def harness(self, name, in_fn, ins, dins, douts):
                         )
 
                         if self.revprimal:
-                            primals, grads = rev_enzyme(*douts_backend, *ins_backend)
+                            primals, grads = rev_enzyme(*adout, *ins_backend)
                         else:
-                            grads = rev_enzyme(*douts_backend, *ins_backend)
+                            grads = rev_enzyme(*adout, *ins_backend)
                             assert grads is not None
 
                         if self.revprimal and primres is not None:
@@ -355,7 +375,7 @@ def harness(self, name, in_fn, ins, dins, douts):
                             sep="\t",
                         )
 
-                    if pipeline is None or pipeline.mlir_ad():
+                    if pipeline is None or (pipeline.mlir_ad() and self.mlirad_rev):
                         rfn_enzyme = (
                             in_fn
                             if pipeline is None
@@ -375,9 +395,9 @@ def harness(self, name, in_fn, ins, dins, douts):
                         )
 
                         if self.revprimal:
-                            primals, grads = rev_enzyme(*douts_backend, *ins_backend)
+                            primals, grads = rev_enzyme(*adout, *ins_backend)
                         else:
-                            grads = rev_enzyme(*douts_backend, *ins_backend)
+                            grads = rev_enzyme(*adout, *ins_backend)
                             assert grads is not None
 
                         if self.revprimal and primres is not None:
@@ -400,4 +420,3 @@ def harness(self, name, in_fn, ins, dins, douts):
                             timeit(csv_filename, pname, "BothRev", revstr, {'rev': rev_enzyme} | revins, self.count),
                             sep="\t",
                         )
-            assert revres is not None
diff --git a/test/vit.py b/test/vit.py
new file mode 100644
index 00000000..66acefc2
--- /dev/null
+++ b/test/vit.py
@@ -0,0 +1,81 @@
+from transformers import AutoImageProcessor, FlaxViTForImageClassification, ViTConfig
+from PIL import Image
+import jax
+import requests
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import absltest
+import jax.random
+from enzyme_ad.jax import JaXPipeline, hlo_opts
+from test_utils import *
+import llama
+
+# Define the pipelines as per your framework
+pipelines = [
+    ("JaX", None, CurBackends),
+    ("JaXPipe", JaXPipeline(), CurBackends),
+    (
+        "HLOOpt",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "canonicalize,cse,enzyme-hlo-opt,cse"
+        ),
+        CurBackends,
+    ),
+    ("PartOpt", JaXPipeline(llama.partialopt), CurBackends),
+    ("DefOpt", JaXPipeline(hlo_opts()), CurBackends),
+    (
+        "EqSat",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "equality-saturation-pass"
+        ),
+        CurBackends,
+    ),
+]
+
+# Load the image processor and model
+processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
+model = FlaxViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
+
+# Define the ViT forward pass
+def vit_forward(pixel_values):
+    """ Vision Transformer forward pass """
+    outputs = model(pixel_values=pixel_values)
+    return outputs.logits
+
+# Define the test class for ViT using absltest
+class VisionTransformerTest(EnzymeJaxTest):
+    def setUp(self):
+        # Create a ViTConfig object instead of using a dictionary
+        self.config = ViTConfig.from_pretrained("google/vit-base-patch16-224")
+
+        # Sample image from COCO dataset
+        url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        pil_image = Image.open(requests.get(url, stream=True).raw)
+        np_image = np.array(pil_image)
+
+        # Process the image
+        inputs = processor(images=np_image, return_tensors="np")
+        pixel_values = inputs["pixel_values"]  # Extract the actual NumPy array
+
+        # Run the ViT forward pass with the processed input
+        self.logits = vit_forward(pixel_values)
+
+        # Setup for the test harness
+        self.fn = vit_forward
+        self.name = "vit"
+        self.count = 10
+        self.revprimal = False
+        self.AllPipelines = pipelines
+        self.AllBackends = CurBackends
+
+        # Input and output setup
+        self.ins = [pixel_values]
+        self.dins = [pixel_values]
+        self.douts = [self.logits]
+        self.tol = 5e-5
+
+# Running the tests with absltest
+if __name__ == "__main__":
+    absltest.main()