requirements.txt

# This file was autogenerated by uv via the following command:
#    uv export --no-hashes --frozen --output-file=requirements.txt
-e .
absl-py==2.4.0
    # via
    #   chex
    #   distrax
    #   dm-tree
    #   kfac-jax
    #   optax
    #   orbax-checkpoint
    #   orbax-export
    #   tfp-nightly
aiofiles==25.1.0
    # via orbax-checkpoint
attrs==26.1.0
    # via dm-tree
beartype==0.22.9
    # via
    #   plum-dispatch
    #   pyserde
casefy==1.1.0
    # via pyserde
chex==0.1.91
    # via distrax
click==8.3.2
    # via jaqmc
cloudpickle==3.1.2
    # via tfp-nightly
colorama==0.4.6 ; sys_platform == 'win32'
    # via
    #   click
    #   pytest
coverage==7.13.5
    # via pytest-cov
dataclasses-json==0.6.7
    # via orbax-export
decorator==5.2.1
    # via tfp-nightly
distrax==0.1.7
    # via kfac-jax
dm-tree==0.1.10
    # via
    #   kfac-jax
    #   tfp-nightly
etils==1.14.0
    # via
    #   orbax-checkpoint
    #   orbax-export
execnet==2.1.2
    # via pytest-xdist
flax==0.12.6
    # via jaqmc
folx @ git+https://github.com/microsoft/folx@47a0e8d97a311686115aae353a85a45ae6432274
    # via jaqmc
fsspec==2026.3.0
    # via
    #   etils
    #   jaqmc
    #   universal-pathlib
gast==0.7.0
    # via tfp-nightly
h5py==3.16.0
    # via
    #   jaqmc
    #   pyscf
humanize==4.15.0
    # via orbax-checkpoint
immutabledict==4.3.1
    # via kfac-jax
iniconfig==2.3.0
    # via pytest
jax==0.9.1
    # via
    #   chex
    #   distrax
    #   flax
    #   folx
    #   jaqmc
    #   kfac-jax
    #   optax
    #   orbax-checkpoint
    #   orbax-export
jaxlib==0.9.1
    # via
    #   chex
    #   distrax
    #   jax
    #   kfac-jax
    #   optax
    #   orbax-export
jaxtyping==0.3.9
    # via
    #   folx
    #   orbax-export
jinja2==3.1.6
    # via pyserde
kfac-jax==0.0.8
    # via jaqmc
librt==0.9.0 ; platform_python_implementation != 'PyPy'
    # via mypy
markdown-it-py==4.0.0
    # via rich
markupsafe==3.0.3
    # via jinja2
marshmallow==3.26.2
    # via dataclasses-json
mdurl==0.1.2
    # via markdown-it-py
ml-dtypes==0.5.4
    # via
    #   jax
    #   jaxlib
    #   tensorstore
msgpack==1.1.2
    # via
    #   flax
    #   orbax-checkpoint
mypy==1.20.0
mypy-extensions==1.1.0
    # via
    #   mypy
    #   typing-inspect
nest-asyncio==1.6.0 ; sys_platform == 'win32'
    # via orbax-checkpoint
numpy==2.4.4
    # via
    #   chex
    #   distrax
    #   dm-tree
    #   flax
    #   folx
    #   h5py
    #   jaqmc
    #   jax
    #   jaxlib
    #   kfac-jax
    #   ml-dtypes
    #   optax
    #   orbax-checkpoint
    #   orbax-export
    #   pandas
    #   pyscf
    #   scipy
    #   tensorstore
    #   tfp-nightly
    #   treescope
opt-einsum==3.4.0
    # via jax
optax==0.2.8
    # via
    #   flax
    #   jaqmc
    #   kfac-jax
orbax-checkpoint==0.11.34
    # via
    #   flax
    #   orbax-export
orbax-export==0.0.8
    # via flax
packaging==26.0
    # via
    #   marshmallow
    #   pytest
pandas==3.0.2
    # via jaqmc
pathlib-abc==0.5.2
    # via universal-pathlib
pathspec==1.0.4
    # via mypy
pluggy==1.6.0
    # via
    #   pytest
    #   pytest-cov
plum-dispatch==2.8.0
    # via pyserde
protobuf==7.34.1
    # via
    #   orbax-checkpoint
    #   orbax-export
psutil==7.2.2
    # via orbax-checkpoint
pygments==2.20.0
    # via
    #   jaqmc
    #   pytest
    #   rich
pyscf==2.12.1
    # via jaqmc
pyserde==0.31.2
    # via jaqmc
pytest==9.0.3
    # via
    #   pytest-cov
    #   pytest-mock
    #   pytest-xdist
pytest-cov==7.1.0
pytest-mock==3.15.1
pytest-xdist==3.8.0
python-dateutil==2.9.0.post0
    # via pandas
pyyaml==6.0.3
    # via
    #   flax
    #   orbax-checkpoint
rich==14.3.3
    # via
    #   flax
    #   plum-dispatch
ruff==0.15.10
scipy==1.17.1
    # via
    #   jax
    #   jaxlib
    #   pyscf
setuptools==82.0.1
    # via pyscf
simplejson==3.20.2
    # via orbax-checkpoint
six==1.17.0
    # via
    #   python-dateutil
    #   tfp-nightly
tensorstore==0.1.82
    # via
    #   flax
    #   orbax-checkpoint
tfp-nightly==0.26.0.dev20260410
    # via distrax
toolz==1.1.0
    # via chex
treescope==0.1.10
    # via flax
types-pyyaml==6.0.12.20260408
typing-extensions==4.15.0
    # via
    #   chex
    #   etils
    #   flax
    #   kfac-jax
    #   mypy
    #   orbax-checkpoint
    #   plum-dispatch
    #   pyserde
    #   typing-inspect
typing-inspect==0.9.0
    # via
    #   dataclasses-json
    #   pyserde
tzdata==2026.1 ; sys_platform == 'emscripten' or sys_platform == 'win32'
    # via pandas
universal-pathlib==0.3.10
    # via jaqmc
uvloop==0.22.1 ; sys_platform != 'win32'
    # via orbax-checkpoint
wadler-lindig==0.1.7
    # via jaxtyping
wrapt==2.1.2
    # via dm-tree
zipp==3.23.0
    # via etils