├── .github ├── FUNDING.yml └── workflows │ ├── create-release.yml │ ├── publish-package.yml │ └── run-tests.yml ├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── poetry.lock ├── pyproject.toml ├── scripts └── update_version.py ├── simple_pytree ├── __init__.py ├── dataclass.py └── pytree.py └── tests └── test_pytree.py /.github/FUNDING.yml: -------------------------------------------------------------------------------- 1 | # These are supported funding model platforms 2 | 3 | github: [cgarciae] 4 | # patreon: # Replace with a single Patreon username 5 | # open_collective: # Replace with a single Open Collective username 6 | # ko_fi: # Replace with a single Ko-fi username 7 | # tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel 8 | # community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry 9 | # liberapay: # Replace with a single Liberapay username 10 | # issuehunt: # Replace with a single IssueHunt username 11 | # otechie: # Replace with a single Otechie username 12 | # lfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry 13 | # custom: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2'] 14 | -------------------------------------------------------------------------------- /.github/workflows/create-release.yml: -------------------------------------------------------------------------------- 1 | name: Create Release 2 | 3 | on: 4 | create 5 | 6 | jobs: 7 | create-release: 8 | if: startsWith(github.ref_name, 'version-') && endsWith(github.ref_name, '-create-release') 9 | name: Create Release 10 | runs-on: ubuntu-latest 11 | steps: 12 | - name: Check out the code 13 | uses: actions/checkout@v3 14 | 15 | - name: Set up Python 3.8 16 | uses: actions/setup-python@v4 17 | with: 18 | python-version: 3.8 19 | 20 | - name: Setup 21 | id: setup 22 | run: | 23 | # install python dependencies 24 | pip install typer==0.4.0 click==8.0.3 25 | 26 | # variables 27 | RELEASE_VERSION='${{ github.ref_name }}' 28 | RELEASE_VERSION=${RELEASE_VERSION//version-/} 29 | RELEASE_VERSION=${RELEASE_VERSION//-create-release/} 30 | echo "::set-output name=RELEASE_VERSION::${RELEASE_VERSION}" 31 | 32 | 33 | - name: Update version 34 | run: | 35 | RELEASE_VERSION='${{ steps.setup.outputs.RELEASE_VERSION }}' 36 | 37 | # setup git 38 | git config --local user.email "github-actions[bot]@users.noreply.github.com" 39 | git config --local user.name "github-actions[bot]" 40 | 41 | # switch to main 42 | git pull origin main 43 | git checkout main 44 | 45 | # update version 46 | python scripts/update_version.py $RELEASE_VERSION 47 | git commit -am "Update version to $RELEASE_VERSION" 48 | 49 | # create tag 50 | git fetch --tags 51 | git tag $RELEASE_VERSION 52 | 53 | # push to main 54 | git push 55 | git push --tags 56 | 57 | # delete branch 58 | git push -d origin ${{ github.ref_name }} 59 | 60 | - name: Create Release 61 | uses: actions/create-release@v1 62 | with: 63 | tag_name: ${{ steps.setup.outputs.RELEASE_VERSION }} 64 | release_name: ${{ steps.setup.outputs.RELEASE_VERSION }} 65 | draft: true 66 | env: 67 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} -------------------------------------------------------------------------------- /.github/workflows/publish-package.yml: -------------------------------------------------------------------------------- 1 | name: Publish Package 2 | on: 3 | release: 4 | types: [published] 5 | jobs: 6 | publish-docs-and-package: 7 | name: Publish Docs and Package 8 | runs-on: ubuntu-latest 9 | steps: 10 | - name: Check out the code 11 | uses: actions/checkout@v3 12 | 13 | - name: Set up Python 3.8 14 | uses: actions/setup-python@v4 15 | with: 16 | python-version: 3.8 17 | 18 | - name: Install Poetry 19 | uses: snok/install-poetry@v1.3.3 20 | with: 21 | version: 1.4.0 22 | 23 | - name: Setup Poetry 24 | run: | 25 | poetry config virtualenvs.in-project true 26 | 27 | - name: Cache 28 | id: cache 29 | uses: actions/cache@v3.2.2 30 | with: 31 | path: '.venv' 32 | key: publish-package-${{ hashFiles('poetry.lock') }} 33 | 34 | - name: Install Dependencies 35 | if: steps.cache.outputs.cache-hit != 'true' 36 | run: | 37 | poetry install --without dev 38 | 39 | - name: Install Package 40 | run: | 41 | poetry install --without dev 42 | 43 | # ---------------------------------------- 44 | # No docs for now 45 | # ---------------------------------------- 46 | # - name: Build Docs 🔨 47 | # run: | 48 | # cp README.md docs/index.md 49 | # poetry run mkdocs build 50 | 51 | # - name: Deploy Page 🚀 52 | # uses: JamesIves/github-pages-deploy-action@4.1.6 53 | # with: 54 | # branch: gh-pages 55 | # folder: site 56 | 57 | - name: Publish to PyPI 58 | run: | 59 | poetry build 60 | poetry publish \ 61 | --username ${{ secrets.PYPI_USERNAME }} \ 62 | --password ${{ secrets.PYPI_PASSWORD }} 63 | -------------------------------------------------------------------------------- /.github/workflows/run-tests.yml: -------------------------------------------------------------------------------- 1 | # Checks that we can build and validate the Unittest 2 | name: Run Tests 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | jobs: 9 | pre-commit: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v3 13 | - uses: actions/setup-python@v4 14 | - uses: pre-commit/action@v3.0.0 15 | test: 16 | name: Run Tests 17 | runs-on: ubuntu-latest 18 | strategy: 19 | matrix: 20 | python-version: ['3.8', '3.9', '3.10'] 21 | steps: 22 | - name: Check out the code 23 | uses: actions/checkout@v3 24 | with: 25 | fetch-depth: 1 26 | - name: Set up Python ${{ matrix.python-version }} 27 | uses: actions/setup-python@v4 28 | with: 29 | python-version: ${{ matrix.python-version }} 30 | 31 | - name: Install Poetry 32 | uses: snok/install-poetry@v1.3.3 33 | with: 34 | version: 1.4.0 35 | 36 | - name: Setup Poetry 37 | run: | 38 | poetry config virtualenvs.in-project true 39 | 40 | - name: Cache 41 | id: cache 42 | uses: actions/cache@v3.2.2 43 | with: 44 | path: '.venv' 45 | key: run-tests-${{ hashFiles('poetry.lock') }} 46 | 47 | - name: Install Dependencies 48 | if: steps.cache.outputs.cache-hit != 'true' 49 | run: | 50 | if [ -d ".venv" ]; then rm -rf .venv; fi 51 | poetry run pip install -U certifi 52 | poetry install 53 | 54 | - name: Install Package 55 | run: | 56 | poetry install 57 | 58 | - name: Run Tests 59 | run: poetry run pytest --cov=simple_pytree --cov-report=term-missing --cov-report=xml 60 | 61 | - name: Upload coverage 62 | uses: codecov/codecov-action@v3 63 | 64 | test-import: 65 | name: Test Import without Dev Dependencies 66 | if: ${{ !contains(github.event.pull_request.title, 'WIP') }} 67 | runs-on: ubuntu-latest 68 | strategy: 69 | matrix: 70 | python-version: ['3.8', '3.9', '3.10'] 71 | steps: 72 | - name: Check out the code 73 | uses: actions/checkout@v3 74 | with: 75 | fetch-depth: 1 76 | - name: Set up Python ${{ matrix.python-version }} 77 | uses: actions/setup-python@v4 78 | with: 79 | python-version: ${{ matrix.python-version }} 80 | 81 | - name: Install Poetry 82 | uses: snok/install-poetry@v1.3.3 83 | with: 84 | version: 1.4.0 85 | 86 | - name: Setup Poetry 87 | run: | 88 | poetry config virtualenvs.in-project true 89 | 90 | - name: Cache 91 | id: cache 92 | uses: actions/cache@v3.2.2 93 | with: 94 | path: '.venv' 95 | key: test-import-${{ hashFiles('poetry.lock') }} 96 | 97 | - name: Install Dependencies 98 | if: steps.cache.outputs.cache-hit != 'true' 99 | run: | 100 | poetry run pip install -U certifi 101 | poetry install --only main 102 | 103 | - name: Install Package 104 | run: | 105 | poetry install --only main 106 | 107 | - name: Test Import 108 | run: | 109 | poetry run python -c "import simple_pytree" 110 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # project specific 132 | .vscode 133 | /tmp -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | 2 | repos: 3 | - repo: https://github.com/psf/black 4 | rev: 23.1.0 5 | hooks: 6 | - id: black 7 | args: ["--target-version", "py38", "--fast"] 8 | - repo: https://github.com/pycqa/isort 9 | rev: 5.12.0 10 | hooks: 11 | - id: isort 12 | args: ["--profile", "black"] 13 | 14 | 15 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Cristian Garcia 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | [![codecov](https://codecov.io/gh/cgarciae/simple-pytree/branch/main/graph/badge.svg?token=3IKEUAU3C8)](https://codecov.io/gh/cgarciae/simple-pytree) 4 | 5 | 6 | # Simple Pytree 7 | 8 | A _dead simple_ Python package for creating custom JAX pytree objects. 9 | 10 | * Strives to be minimal, the implementation is just ~100 lines of code 11 | * Has no dependencies other than JAX 12 | * Its compatible with both `dataclasses` and regular classes 13 | * It has no intention of supporting Neural Network use cases (e.g. partitioning) 14 | 15 | ## Installation 16 | 17 | ```bash 18 | pip install simple-pytree 19 | ``` 20 | 21 | ## Usage 22 | 23 | ```python 24 | import jax 25 | from simple_pytree import Pytree 26 | 27 | class Foo(Pytree): 28 | def __init__(self, x, y): 29 | self.x = x 30 | self.y = y 31 | 32 | foo = Foo(1, 2) 33 | foo = jax.tree_map(lambda x: -x, foo) 34 | 35 | assert foo.x == -1 and foo.y == -2 36 | ``` 37 | 38 | ### Static fields 39 | You can mark fields as static by assigning `static_field()` to a class attribute with the same name 40 | as the instance attribute: 41 | 42 | ```python 43 | import jax 44 | from simple_pytree import Pytree, static_field 45 | 46 | class Foo(Pytree): 47 | y = static_field() 48 | 49 | def __init__(self, x, y): 50 | self.x = x 51 | self.y = y 52 | 53 | foo = Foo(1, 2) 54 | foo = jax.tree_map(lambda x: -x, foo) # y is not modified 55 | 56 | assert foo.x == -1 and foo.y == 2 57 | ``` 58 | 59 | Static fields are not included in the pytree leaves, they 60 | are passed as pytree metadata instead. 61 | 62 | ### Dataclasses 63 | `simple_pytree` provides a `dataclass` decorator you can use with classes 64 | that contain `static_field`s: 65 | 66 | ```python 67 | import jax 68 | from simple_pytree import Pytree, dataclass, static_field 69 | 70 | @dataclass 71 | class Foo(Pytree): 72 | x: int 73 | y: int = static_field(default=2) 74 | 75 | foo = Foo(1) 76 | foo = jax.tree_map(lambda x: -x, foo) # y is not modified 77 | 78 | assert foo.x == -1 and foo.y == 2 79 | ``` 80 | `simple_pytree.dataclass` is just a wrapper around `dataclasses.dataclass` but 81 | when used static analysis tools and IDEs will understand that `static_field` is a 82 | field specifier just like `dataclasses.field`. 83 | 84 | ### Mutability 85 | `Pytree` objects are immutable by default after `__init__`: 86 | 87 | ```python 88 | from simple_pytree import Pytree, static_field 89 | 90 | class Foo(Pytree): 91 | y = static_field() 92 | 93 | def __init__(self, x, y): 94 | self.x = x 95 | self.y = y 96 | 97 | foo = Foo(1, 2) 98 | foo.x = 3 # AttributeError 99 | ``` 100 | If you want to make them mutable, you can use the `mutable` argument in class definition: 101 | 102 | ```python 103 | from simple_pytree import Pytree, static_field 104 | 105 | class Foo(Pytree, mutable=True): 106 | y = static_field() 107 | 108 | def __init__(self, x, y): 109 | self.x = x 110 | self.y = y 111 | 112 | foo = Foo(1, 2) 113 | foo.x = 3 # OK 114 | ``` 115 | 116 | ### Replacing fields 117 | 118 | If you want to make a copy of a `Pytree` object with some fields modified, you can use the `.replace()` method: 119 | 120 | ```python 121 | from simple_pytree import Pytree, static_field 122 | 123 | class Foo(Pytree): 124 | y = static_field() 125 | 126 | def __init__(self, x, y): 127 | self.x = x 128 | self.y = y 129 | 130 | foo = Foo(1, 2) 131 | foo = foo.replace(x=10) 132 | 133 | assert foo.x == 10 and foo.y == 2 134 | ``` 135 | 136 | `replace` works for both mutable and immutable `Pytree` objects. If the class 137 | is a `dataclass`, `replace` internally use `dataclasses.replace`. 138 | 139 | -------------------------------------------------------------------------------- /poetry.lock: -------------------------------------------------------------------------------- 1 | # This file is automatically @generated by Poetry 1.4.0 and should not be changed by hand. 2 | 3 | [[package]] 4 | name = "absl-py" 5 | version = "1.4.0" 6 | description = "Abseil Python Common Libraries, see https://github.com/abseil/abseil-py." 7 | category = "dev" 8 | optional = false 9 | python-versions = ">=3.6" 10 | files = [ 11 | {file = "absl-py-1.4.0.tar.gz", hash = "sha256:d2c244d01048ba476e7c080bd2c6df5e141d211de80223460d5b3b8a2a58433d"}, 12 | {file = "absl_py-1.4.0-py3-none-any.whl", hash = "sha256:0d3fe606adfa4f7db64792dd4c7aee4ee0c38ab75dfd353b7a83ed3e957fcb47"}, 13 | ] 14 | 15 | [[package]] 16 | name = "black" 17 | version = "23.1.0" 18 | description = "The uncompromising code formatter." 19 | category = "dev" 20 | optional = false 21 | python-versions = ">=3.7" 22 | files = [ 23 | {file = "black-23.1.0-cp310-cp310-macosx_10_16_arm64.whl", hash = "sha256:b6a92a41ee34b883b359998f0c8e6eb8e99803aa8bf3123bf2b2e6fec505a221"}, 24 | {file = "black-23.1.0-cp310-cp310-macosx_10_16_universal2.whl", hash = "sha256:57c18c5165c1dbe291d5306e53fb3988122890e57bd9b3dcb75f967f13411a26"}, 25 | {file = "black-23.1.0-cp310-cp310-macosx_10_16_x86_64.whl", hash = "sha256:9880d7d419bb7e709b37e28deb5e68a49227713b623c72b2b931028ea65f619b"}, 26 | {file = "black-23.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e6663f91b6feca5d06f2ccd49a10f254f9298cc1f7f49c46e498a0771b507104"}, 27 | {file = "black-23.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:9afd3f493666a0cd8f8df9a0200c6359ac53940cbde049dcb1a7eb6ee2dd7074"}, 28 | {file = "black-23.1.0-cp311-cp311-macosx_10_16_arm64.whl", hash = "sha256:bfffba28dc52a58f04492181392ee380e95262af14ee01d4bc7bb1b1c6ca8d27"}, 29 | {file = "black-23.1.0-cp311-cp311-macosx_10_16_universal2.whl", hash = "sha256:c1c476bc7b7d021321e7d93dc2cbd78ce103b84d5a4cf97ed535fbc0d6660648"}, 30 | {file = "black-23.1.0-cp311-cp311-macosx_10_16_x86_64.whl", hash = "sha256:382998821f58e5c8238d3166c492139573325287820963d2f7de4d518bd76958"}, 31 | {file = "black-23.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2bf649fda611c8550ca9d7592b69f0637218c2369b7744694c5e4902873b2f3a"}, 32 | {file = "black-23.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:121ca7f10b4a01fd99951234abdbd97728e1240be89fde18480ffac16503d481"}, 33 | {file = "black-23.1.0-cp37-cp37m-macosx_10_16_x86_64.whl", hash = "sha256:a8471939da5e824b891b25751955be52ee7f8a30a916d570a5ba8e0f2eb2ecad"}, 34 | {file = "black-23.1.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8178318cb74f98bc571eef19068f6ab5613b3e59d4f47771582f04e175570ed8"}, 35 | {file = "black-23.1.0-cp37-cp37m-win_amd64.whl", hash = "sha256:a436e7881d33acaf2536c46a454bb964a50eff59b21b51c6ccf5a40601fbef24"}, 36 | {file = "black-23.1.0-cp38-cp38-macosx_10_16_arm64.whl", hash = "sha256:a59db0a2094d2259c554676403fa2fac3473ccf1354c1c63eccf7ae65aac8ab6"}, 37 | {file = "black-23.1.0-cp38-cp38-macosx_10_16_universal2.whl", hash = "sha256:0052dba51dec07ed029ed61b18183942043e00008ec65d5028814afaab9a22fd"}, 38 | {file = "black-23.1.0-cp38-cp38-macosx_10_16_x86_64.whl", hash = "sha256:49f7b39e30f326a34b5c9a4213213a6b221d7ae9d58ec70df1c4a307cf2a1580"}, 39 | {file = "black-23.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:162e37d49e93bd6eb6f1afc3e17a3d23a823042530c37c3c42eeeaf026f38468"}, 40 | {file = "black-23.1.0-cp38-cp38-win_amd64.whl", hash = "sha256:8b70eb40a78dfac24842458476135f9b99ab952dd3f2dab738c1881a9b38b753"}, 41 | {file = "black-23.1.0-cp39-cp39-macosx_10_16_arm64.whl", hash = "sha256:a29650759a6a0944e7cca036674655c2f0f63806ddecc45ed40b7b8aa314b651"}, 42 | {file = "black-23.1.0-cp39-cp39-macosx_10_16_universal2.whl", hash = "sha256:bb460c8561c8c1bec7824ecbc3ce085eb50005883a6203dcfb0122e95797ee06"}, 43 | {file = "black-23.1.0-cp39-cp39-macosx_10_16_x86_64.whl", hash = "sha256:c91dfc2c2a4e50df0026f88d2215e166616e0c80e86004d0003ece0488db2739"}, 44 | {file = "black-23.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2a951cc83ab535d248c89f300eccbd625e80ab880fbcfb5ac8afb5f01a258ac9"}, 45 | {file = "black-23.1.0-cp39-cp39-win_amd64.whl", hash = "sha256:0680d4380db3719ebcfb2613f34e86c8e6d15ffeabcf8ec59355c5e7b85bb555"}, 46 | {file = "black-23.1.0-py3-none-any.whl", hash = "sha256:7a0f701d314cfa0896b9001df70a530eb2472babb76086344e688829efd97d32"}, 47 | {file = "black-23.1.0.tar.gz", hash = "sha256:b0bd97bea8903f5a2ba7219257a44e3f1f9d00073d6cc1add68f0beec69692ac"}, 48 | ] 49 | 50 | [package.dependencies] 51 | click = ">=8.0.0" 52 | mypy-extensions = ">=0.4.3" 53 | packaging = ">=22.0" 54 | pathspec = ">=0.9.0" 55 | platformdirs = ">=2" 56 | tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""} 57 | typing-extensions = {version = ">=3.10.0.0", markers = "python_version < \"3.10\""} 58 | 59 | [package.extras] 60 | colorama = ["colorama (>=0.4.3)"] 61 | d = ["aiohttp (>=3.7.4)"] 62 | jupyter = ["ipython (>=7.8.0)", "tokenize-rt (>=3.2.0)"] 63 | uvloop = ["uvloop (>=0.15.2)"] 64 | 65 | [[package]] 66 | name = "cached-property" 67 | version = "1.5.2" 68 | description = "A decorator for caching properties in classes." 69 | category = "dev" 70 | optional = false 71 | python-versions = "*" 72 | files = [ 73 | {file = "cached-property-1.5.2.tar.gz", hash = "sha256:9fa5755838eecbb2d234c3aa390bd80fbd3ac6b6869109bfc1b499f7bd89a130"}, 74 | {file = "cached_property-1.5.2-py2.py3-none-any.whl", hash = "sha256:df4f613cf7ad9a588cc381aaf4a512d26265ecebd5eb9e1ba12f1319eb85a6a0"}, 75 | ] 76 | 77 | [[package]] 78 | name = "cfgv" 79 | version = "3.3.1" 80 | description = "Validate configuration and produce human readable error messages." 81 | category = "dev" 82 | optional = false 83 | python-versions = ">=3.6.1" 84 | files = [ 85 | {file = "cfgv-3.3.1-py2.py3-none-any.whl", hash = "sha256:c6a0883f3917a037485059700b9e75da2464e6c27051014ad85ba6aaa5884426"}, 86 | {file = "cfgv-3.3.1.tar.gz", hash = "sha256:f5a830efb9ce7a445376bb66ec94c638a9787422f96264c98edc6bdeed8ab736"}, 87 | ] 88 | 89 | [[package]] 90 | name = "chex" 91 | version = "0.1.7" 92 | description = "Chex: Testing made fun, in JAX!" 93 | category = "dev" 94 | optional = false 95 | python-versions = ">=3.8" 96 | files = [ 97 | {file = "chex-0.1.7-py3-none-any.whl", hash = "sha256:9f583015303b1205443843c0b55849bb287f1dfdbd22d9907b1ebb04f964d93e"}, 98 | {file = "chex-0.1.7.tar.gz", hash = "sha256:74ed49799ac4d229881456d468136f1b19a9f9839e3de72b058824e2a4f4dedd"}, 99 | ] 100 | 101 | [package.dependencies] 102 | absl-py = ">=0.9.0" 103 | dm-tree = ">=0.1.5" 104 | jax = ">=0.4.6" 105 | jaxlib = ">=0.1.37" 106 | numpy = ">=1.18.0" 107 | toolz = ">=0.9.0" 108 | typing-extensions = {version = ">=4.2.0", markers = "python_version < \"3.11\""} 109 | 110 | [[package]] 111 | name = "click" 112 | version = "8.1.3" 113 | description = "Composable command line interface toolkit" 114 | category = "dev" 115 | optional = false 116 | python-versions = ">=3.7" 117 | files = [ 118 | {file = "click-8.1.3-py3-none-any.whl", hash = "sha256:bb4d8133cb15a609f44e8213d9b391b0809795062913b383c62be0ee95b1db48"}, 119 | {file = "click-8.1.3.tar.gz", hash = "sha256:7682dc8afb30297001674575ea00d1814d808d6a36af415a82bd481d37ba7b8e"}, 120 | ] 121 | 122 | [package.dependencies] 123 | colorama = {version = "*", markers = "platform_system == \"Windows\""} 124 | 125 | [[package]] 126 | name = "colorama" 127 | version = "0.4.6" 128 | description = "Cross-platform colored terminal text." 129 | category = "dev" 130 | optional = false 131 | python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" 132 | files = [ 133 | {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, 134 | {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, 135 | ] 136 | 137 | [[package]] 138 | name = "coverage" 139 | version = "7.2.3" 140 | description = "Code coverage measurement for Python" 141 | category = "dev" 142 | optional = false 143 | python-versions = ">=3.7" 144 | files = [ 145 | {file = "coverage-7.2.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e58c0d41d336569d63d1b113bd573db8363bc4146f39444125b7f8060e4e04f5"}, 146 | {file = "coverage-7.2.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:344e714bd0fe921fc72d97404ebbdbf9127bac0ca1ff66d7b79efc143cf7c0c4"}, 147 | {file = "coverage-7.2.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:974bc90d6f6c1e59ceb1516ab00cf1cdfbb2e555795d49fa9571d611f449bcb2"}, 148 | {file = "coverage-7.2.3-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0743b0035d4b0e32bc1df5de70fba3059662ace5b9a2a86a9f894cfe66569013"}, 149 | {file = "coverage-7.2.3-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5d0391fb4cfc171ce40437f67eb050a340fdbd0f9f49d6353a387f1b7f9dd4fa"}, 150 | {file = "coverage-7.2.3-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:4a42e1eff0ca9a7cb7dc9ecda41dfc7cbc17cb1d02117214be0561bd1134772b"}, 151 | {file = "coverage-7.2.3-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:be19931a8dcbe6ab464f3339966856996b12a00f9fe53f346ab3be872d03e257"}, 152 | {file = "coverage-7.2.3-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:72fcae5bcac3333a4cf3b8f34eec99cea1187acd55af723bcbd559adfdcb5535"}, 153 | {file = "coverage-7.2.3-cp310-cp310-win32.whl", hash = "sha256:aeae2aa38395b18106e552833f2a50c27ea0000122bde421c31d11ed7e6f9c91"}, 154 | {file = "coverage-7.2.3-cp310-cp310-win_amd64.whl", hash = "sha256:83957d349838a636e768251c7e9979e899a569794b44c3728eaebd11d848e58e"}, 155 | {file = "coverage-7.2.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:dfd393094cd82ceb9b40df4c77976015a314b267d498268a076e940fe7be6b79"}, 156 | {file = "coverage-7.2.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:182eb9ac3f2b4874a1f41b78b87db20b66da6b9cdc32737fbbf4fea0c35b23fc"}, 157 | {file = "coverage-7.2.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1bb1e77a9a311346294621be905ea8a2c30d3ad371fc15bb72e98bfcfae532df"}, 158 | {file = "coverage-7.2.3-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ca0f34363e2634deffd390a0fef1aa99168ae9ed2af01af4a1f5865e362f8623"}, 159 | {file = "coverage-7.2.3-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:55416d7385774285b6e2a5feca0af9652f7f444a4fa3d29d8ab052fafef9d00d"}, 160 | {file = "coverage-7.2.3-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:06ddd9c0249a0546997fdda5a30fbcb40f23926df0a874a60a8a185bc3a87d93"}, 161 | {file = "coverage-7.2.3-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:fff5aaa6becf2c6a1699ae6a39e2e6fb0672c2d42eca8eb0cafa91cf2e9bd312"}, 162 | {file = "coverage-7.2.3-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:ea53151d87c52e98133eb8ac78f1206498c015849662ca8dc246255265d9c3c4"}, 163 | {file = "coverage-7.2.3-cp311-cp311-win32.whl", hash = "sha256:8f6c930fd70d91ddee53194e93029e3ef2aabe26725aa3c2753df057e296b925"}, 164 | {file = "coverage-7.2.3-cp311-cp311-win_amd64.whl", hash = "sha256:fa546d66639d69aa967bf08156eb8c9d0cd6f6de84be9e8c9819f52ad499c910"}, 165 | {file = "coverage-7.2.3-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:b2317d5ed777bf5a033e83d4f1389fd4ef045763141d8f10eb09a7035cee774c"}, 166 | {file = "coverage-7.2.3-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be9824c1c874b73b96288c6d3de793bf7f3a597770205068c6163ea1f326e8b9"}, 167 | {file = "coverage-7.2.3-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2c3b2803e730dc2797a017335827e9da6da0e84c745ce0f552e66400abdfb9a1"}, 168 | {file = "coverage-7.2.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f69770f5ca1994cb32c38965e95f57504d3aea96b6c024624fdd5bb1aa494a1"}, 169 | {file = "coverage-7.2.3-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:1127b16220f7bfb3f1049ed4a62d26d81970a723544e8252db0efde853268e21"}, 170 | {file = "coverage-7.2.3-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:aa784405f0c640940595fa0f14064d8e84aff0b0f762fa18393e2760a2cf5841"}, 171 | {file = "coverage-7.2.3-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:3146b8e16fa60427e03884301bf8209221f5761ac754ee6b267642a2fd354c48"}, 172 | {file = "coverage-7.2.3-cp37-cp37m-win32.whl", hash = "sha256:1fd78b911aea9cec3b7e1e2622c8018d51c0d2bbcf8faaf53c2497eb114911c1"}, 173 | {file = "coverage-7.2.3-cp37-cp37m-win_amd64.whl", hash = "sha256:0f3736a5d34e091b0a611964c6262fd68ca4363df56185902528f0b75dbb9c1f"}, 174 | {file = "coverage-7.2.3-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:981b4df72c93e3bc04478153df516d385317628bd9c10be699c93c26ddcca8ab"}, 175 | {file = "coverage-7.2.3-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:c0045f8f23a5fb30b2eb3b8a83664d8dc4fb58faddf8155d7109166adb9f2040"}, 176 | {file = "coverage-7.2.3-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f760073fcf8f3d6933178d67754f4f2d4e924e321f4bb0dcef0424ca0215eba1"}, 177 | {file = "coverage-7.2.3-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c86bd45d1659b1ae3d0ba1909326b03598affbc9ed71520e0ff8c31a993ad911"}, 178 | {file = "coverage-7.2.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:172db976ae6327ed4728e2507daf8a4de73c7cc89796483e0a9198fd2e47b462"}, 179 | {file = "coverage-7.2.3-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:d2a3a6146fe9319926e1d477842ca2a63fe99af5ae690b1f5c11e6af074a6b5c"}, 180 | {file = "coverage-7.2.3-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:f649dd53833b495c3ebd04d6eec58479454a1784987af8afb77540d6c1767abd"}, 181 | {file = "coverage-7.2.3-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:7c4ed4e9f3b123aa403ab424430b426a1992e6f4c8fd3cb56ea520446e04d152"}, 182 | {file = "coverage-7.2.3-cp38-cp38-win32.whl", hash = "sha256:eb0edc3ce9760d2f21637766c3aa04822030e7451981ce569a1b3456b7053f22"}, 183 | {file = "coverage-7.2.3-cp38-cp38-win_amd64.whl", hash = "sha256:63cdeaac4ae85a179a8d6bc09b77b564c096250d759eed343a89d91bce8b6367"}, 184 | {file = "coverage-7.2.3-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:20d1a2a76bb4eb00e4d36b9699f9b7aba93271c9c29220ad4c6a9581a0320235"}, 185 | {file = "coverage-7.2.3-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:4ea748802cc0de4de92ef8244dd84ffd793bd2e7be784cd8394d557a3c751e21"}, 186 | {file = "coverage-7.2.3-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:21b154aba06df42e4b96fc915512ab39595105f6c483991287021ed95776d934"}, 187 | {file = "coverage-7.2.3-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:fd214917cabdd6f673a29d708574e9fbdb892cb77eb426d0eae3490d95ca7859"}, 188 | {file = "coverage-7.2.3-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2c2e58e45fe53fab81f85474e5d4d226eeab0f27b45aa062856c89389da2f0d9"}, 189 | {file = "coverage-7.2.3-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:87ecc7c9a1a9f912e306997ffee020297ccb5ea388421fe62a2a02747e4d5539"}, 190 | {file = "coverage-7.2.3-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:387065e420aed3c71b61af7e82c7b6bc1c592f7e3c7a66e9f78dd178699da4fe"}, 191 | {file = "coverage-7.2.3-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:ea3f5bc91d7d457da7d48c7a732beaf79d0c8131df3ab278e6bba6297e23c6c4"}, 192 | {file = "coverage-7.2.3-cp39-cp39-win32.whl", hash = "sha256:ae7863a1d8db6a014b6f2ff9c1582ab1aad55a6d25bac19710a8df68921b6e30"}, 193 | {file = "coverage-7.2.3-cp39-cp39-win_amd64.whl", hash = "sha256:3f04becd4fcda03c0160d0da9c8f0c246bc78f2f7af0feea1ec0930e7c93fa4a"}, 194 | {file = "coverage-7.2.3-pp37.pp38.pp39-none-any.whl", hash = "sha256:965ee3e782c7892befc25575fa171b521d33798132692df428a09efacaffe8d0"}, 195 | {file = "coverage-7.2.3.tar.gz", hash = "sha256:d298c2815fa4891edd9abe5ad6e6cb4207104c7dd9fd13aea3fdebf6f9b91259"}, 196 | ] 197 | 198 | [package.dependencies] 199 | tomli = {version = "*", optional = true, markers = "python_full_version <= \"3.11.0a6\" and extra == \"toml\""} 200 | 201 | [package.extras] 202 | toml = ["tomli"] 203 | 204 | [[package]] 205 | name = "distlib" 206 | version = "0.3.6" 207 | description = "Distribution utilities" 208 | category = "dev" 209 | optional = false 210 | python-versions = "*" 211 | files = [ 212 | {file = "distlib-0.3.6-py2.py3-none-any.whl", hash = "sha256:f35c4b692542ca110de7ef0bea44d73981caeb34ca0b9b6b2e6d7790dda8f80e"}, 213 | {file = "distlib-0.3.6.tar.gz", hash = "sha256:14bad2d9b04d3a36127ac97f30b12a19268f211063d8f8ee4f47108896e11b46"}, 214 | ] 215 | 216 | [[package]] 217 | name = "dm-tree" 218 | version = "0.1.8" 219 | description = "Tree is a library for working with nested data structures." 220 | category = "dev" 221 | optional = false 222 | python-versions = "*" 223 | files = [ 224 | {file = "dm-tree-0.1.8.tar.gz", hash = "sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430"}, 225 | {file = "dm_tree-0.1.8-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:35cc164a79336bfcfafb47e5f297898359123bbd3330c1967f0c4994f9cf9f60"}, 226 | {file = "dm_tree-0.1.8-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:39070ba268c0491af9fe7a58644d99e8b4f2cde6e5884ba3380bddc84ed43d5f"}, 227 | {file = "dm_tree-0.1.8-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2869228d9c619074de501a3c10dc7f07c75422f8fab36ecdcb859b6f1b1ec3ef"}, 228 | {file = "dm_tree-0.1.8-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d20f2faa3672b52e5013f4077117bfb99c4cfc0b445d3bde1584c34032b57436"}, 229 | {file = "dm_tree-0.1.8-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5483dca4d7eb1a0d65fe86d3b6a53ae717face83c1f17e0887b1a4a64ae5c410"}, 230 | {file = "dm_tree-0.1.8-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1d7c26e431fc93cc7e0cba867eb000db6a05f6f2b25af11ac4e9dada88fc5bca"}, 231 | {file = "dm_tree-0.1.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e4d714371bb08839e4e5e29024fc95832d9affe129825ef38836b143028bd144"}, 232 | {file = "dm_tree-0.1.8-cp310-cp310-win_amd64.whl", hash = "sha256:d40fa4106ca6edc66760246a08f500ec0c85ef55c762fb4a363f6ee739ba02ee"}, 233 | {file = "dm_tree-0.1.8-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ad16ceba90a56ec47cf45b21856d14962ac314787975ef786efb5e6e9ca75ec7"}, 234 | {file = "dm_tree-0.1.8-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:803bfc53b4659f447ac694dbd04235f94a73ef7c1fd1e0df7c84ac41e0bc963b"}, 235 | {file = "dm_tree-0.1.8-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:378cc8ad93c5fe3590f405a309980721f021c790ca1bdf9b15bb1d59daec57f5"}, 236 | {file = "dm_tree-0.1.8-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1607ce49aa42f010d1e5e616d92ce899d66835d4d8bea49679582435285515de"}, 237 | {file = "dm_tree-0.1.8-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:343a4a4ebaa127451ff971254a4be4084eb4bdc0b2513c32b46f6f728fd03f9e"}, 238 | {file = "dm_tree-0.1.8-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:fa42a605d099ee7d41ba2b5fb75e21423951fd26e5d50583a00471238fb3021d"}, 239 | {file = "dm_tree-0.1.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:83b7764de0d855338abefc6e3ee9fe40d301668310aa3baea3f778ff051f4393"}, 240 | {file = "dm_tree-0.1.8-cp311-cp311-win_amd64.whl", hash = "sha256:a5d819c38c03f0bb5b3b3703c60e4b170355a0fc6b5819325bf3d4ceb3ae7e80"}, 241 | {file = "dm_tree-0.1.8-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:8c60a7eadab64c2278861f56bca320b2720f163dca9d7558103c3b77f2416571"}, 242 | {file = "dm_tree-0.1.8-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:af4b3d372f2477dcd89a6e717e4a575ca35ccc20cc4454a8a4b6f8838a00672d"}, 243 | {file = "dm_tree-0.1.8-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:de287fabc464b8734be251e46e06aa9aa1001f34198da2b6ce07bd197172b9cb"}, 244 | {file = "dm_tree-0.1.8-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:054b461f8176f4bce7a21f7b1870f873a1ced3bdbe1282c816c550bb43c71fa6"}, 245 | {file = "dm_tree-0.1.8-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2f7915660f59c09068e428613c480150180df1060561fd0d1470684ae7007bd1"}, 246 | {file = "dm_tree-0.1.8-cp37-cp37m-win_amd64.whl", hash = "sha256:b9f89a454e98806b44fe9d40ec9eee61f848388f7e79ac2371a55679bd5a3ac6"}, 247 | {file = "dm_tree-0.1.8-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0e9620ccf06393eb6b613b5e366469304622d4ea96ae6540b28a33840e6c89cf"}, 248 | {file = "dm_tree-0.1.8-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:b095ba4f8ca1ba19350fd53cf1f8f3eb0bd406aa28af64a6dfc86707b32a810a"}, 249 | {file = "dm_tree-0.1.8-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:b9bd9b9ccb59409d33d51d84b7668010c04c2af7d4a371632874c1ca356cff3d"}, 250 | {file = "dm_tree-0.1.8-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0d3172394079a86c3a759179c65f64c48d1a42b89495fcf38976d11cc3bb952c"}, 251 | {file = "dm_tree-0.1.8-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d1612fcaecd79023dbc6a6ae48d51a80beb5c385d6f3f6d71688e57bc8d07de8"}, 252 | {file = "dm_tree-0.1.8-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c5c8c12e3fda754ef6af94161bacdaeda816d941995fac415d6855c6c386af68"}, 253 | {file = "dm_tree-0.1.8-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:694c3654cfd2a81552c08ec66bb5c4a3d48fa292b9a181880fb081c36c5b9134"}, 254 | {file = "dm_tree-0.1.8-cp38-cp38-win_amd64.whl", hash = "sha256:bb2d109f42190225112da899b9f3d46d0d5f26aef501c61e43529fe9322530b5"}, 255 | {file = "dm_tree-0.1.8-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:d16e1f2a073604cfcc09f7131ae8d534674f43c3aef4c25742eae295bc60d04f"}, 256 | {file = "dm_tree-0.1.8-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:250b692fb75f45f02e2f58fbef9ab338904ef334b90557565621fa251df267cf"}, 257 | {file = "dm_tree-0.1.8-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:81fce77f22a302d7a5968aebdf4efafef4def7ce96528719a354e6990dcd49c7"}, 258 | {file = "dm_tree-0.1.8-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f7ac31b9aecccb2c6e1ab29706f6ded3eba0c2c69c770322c9c685929c3d6afb"}, 259 | {file = "dm_tree-0.1.8-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1fe962015b2fe1282892b28ebe962faed53c7f98d942da9a4625cbf27baef913"}, 260 | {file = "dm_tree-0.1.8-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c52cbf4f8b3dbd0beaedf44f69fa85eec5e9dede612e08035e06ada6ec9426"}, 261 | {file = "dm_tree-0.1.8-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:181c35521d480d0365f39300542cb6cd7fd2b77351bb43d7acfda15aef63b317"}, 262 | {file = "dm_tree-0.1.8-cp39-cp39-win_amd64.whl", hash = "sha256:8ed3564abed97c806db122c2d3e1a2b64c74a63debe9903aad795167cc301368"}, 263 | ] 264 | 265 | [[package]] 266 | name = "etils" 267 | version = "1.2.0" 268 | description = "Collection of common python utils" 269 | category = "dev" 270 | optional = false 271 | python-versions = ">=3.8" 272 | files = [ 273 | {file = "etils-1.2.0-py3-none-any.whl", hash = "sha256:c6585069b387fdbeed6a2c571b8bcf312ecdb577c95065461e5fad9ed1973989"}, 274 | {file = "etils-1.2.0.tar.gz", hash = "sha256:29d369e2dcf43960d9ee338330579d04badd606c88f015f4e1a38d3adbe446d8"}, 275 | ] 276 | 277 | [package.extras] 278 | all = ["etils[array-types]", "etils[eapp]", "etils[ecolab]", "etils[edc]", "etils[enp]", "etils[epath]", "etils[epy]", "etils[etqdm]", "etils[etree-dm]", "etils[etree-jax]", "etils[etree-tf]", "etils[etree]"] 279 | array-types = ["etils[enp]"] 280 | dev = ["chex", "optree", "pyink", "pylint (>=2.6.0)", "pytest", "pytest-subtests", "pytest-xdist", "torch"] 281 | eapp = ["absl-py", "etils[epy]", "simple_parsing"] 282 | ecolab = ["etils[enp]", "etils[epy]", "jupyter", "mediapy", "numpy"] 283 | edc = ["etils[epy]", "typing_extensions"] 284 | enp = ["etils[epy]", "numpy"] 285 | epath = ["etils[epy]", "importlib_resources", "typing_extensions", "zipp"] 286 | epy = ["typing_extensions"] 287 | etqdm = ["absl-py", "etils[epy]", "tqdm"] 288 | etree = ["etils[array-types]", "etils[enp]", "etils[epy]", "etils[etqdm]"] 289 | etree-dm = ["dm-tree", "etils[etree]"] 290 | etree-jax = ["etils[etree]", "jax[cpu]"] 291 | etree-tf = ["etils[etree]", "tensorflow"] 292 | lazy-imports = ["etils[ecolab]"] 293 | 294 | [[package]] 295 | name = "exceptiongroup" 296 | version = "1.1.1" 297 | description = "Backport of PEP 654 (exception groups)" 298 | category = "dev" 299 | optional = false 300 | python-versions = ">=3.7" 301 | files = [ 302 | {file = "exceptiongroup-1.1.1-py3-none-any.whl", hash = "sha256:232c37c63e4f682982c8b6459f33a8981039e5fb8756b2074364e5055c498c9e"}, 303 | {file = "exceptiongroup-1.1.1.tar.gz", hash = "sha256:d484c3090ba2889ae2928419117447a14daf3c1231d5e30d0aae34f354f01785"}, 304 | ] 305 | 306 | [package.extras] 307 | test = ["pytest (>=6)"] 308 | 309 | [[package]] 310 | name = "filelock" 311 | version = "3.11.0" 312 | description = "A platform independent file lock." 313 | category = "dev" 314 | optional = false 315 | python-versions = ">=3.7" 316 | files = [ 317 | {file = "filelock-3.11.0-py3-none-any.whl", hash = "sha256:f08a52314748335c6460fc8fe40cd5638b85001225db78c2aa01c8c0db83b318"}, 318 | {file = "filelock-3.11.0.tar.gz", hash = "sha256:3618c0da67adcc0506b015fd11ef7faf1b493f0b40d87728e19986b536890c37"}, 319 | ] 320 | 321 | [package.extras] 322 | docs = ["furo (>=2023.3.27)", "sphinx (>=6.1.3)", "sphinx-autodoc-typehints (>=1.22,!=1.23.4)"] 323 | testing = ["covdefaults (>=2.3)", "coverage (>=7.2.2)", "diff-cover (>=7.5)", "pytest (>=7.2.2)", "pytest-cov (>=4)", "pytest-mock (>=3.10)", "pytest-timeout (>=2.1)"] 324 | 325 | [[package]] 326 | name = "flax" 327 | version = "0.6.8" 328 | description = "Flax: A neural network library for JAX designed for flexibility" 329 | category = "dev" 330 | optional = false 331 | python-versions = "*" 332 | files = [ 333 | {file = "flax-0.6.8-py3-none-any.whl", hash = "sha256:221225804c263e39fe3cc8f754dc4192597cb0f063926b2338ea6563604747ed"}, 334 | {file = "flax-0.6.8.tar.gz", hash = "sha256:bf1f81dd5dfbb10c603490531a86b1174ebbc38e5c5e8116a98115c135194c10"}, 335 | ] 336 | 337 | [package.dependencies] 338 | jax = ">=0.4.2" 339 | msgpack = "*" 340 | numpy = ">=1.12" 341 | optax = "*" 342 | orbax = "*" 343 | PyYAML = ">=5.4.1" 344 | rich = ">=11.1" 345 | tensorstore = "*" 346 | typing-extensions = ">=4.1.1" 347 | 348 | [package.extras] 349 | all = ["matplotlib"] 350 | testing = ["atari-py (==0.2.5)", "clu", "einops", "gym (==0.18.3)", "jaxlib", "jraph (>=0.0.6dev0)", "ml-collections", "mypy", "nbstripout", "opencv-python", "pytest", "pytest-cov", "pytest-custom-exit-code", "pytest-xdist (==1.34.0)", "pytype", "sentencepiece", "tensorflow", "tensorflow-datasets", "tensorflow-text (>=2.11.0)", "torch"] 351 | 352 | [[package]] 353 | name = "identify" 354 | version = "2.5.22" 355 | description = "File identification library for Python" 356 | category = "dev" 357 | optional = false 358 | python-versions = ">=3.7" 359 | files = [ 360 | {file = "identify-2.5.22-py2.py3-none-any.whl", hash = "sha256:f0faad595a4687053669c112004178149f6c326db71ee999ae4636685753ad2f"}, 361 | {file = "identify-2.5.22.tar.gz", hash = "sha256:f7a93d6cf98e29bd07663c60728e7a4057615068d7a639d132dc883b2d54d31e"}, 362 | ] 363 | 364 | [package.extras] 365 | license = ["ukkonen"] 366 | 367 | [[package]] 368 | name = "importlib-resources" 369 | version = "5.12.0" 370 | description = "Read resources from Python packages" 371 | category = "dev" 372 | optional = false 373 | python-versions = ">=3.7" 374 | files = [ 375 | {file = "importlib_resources-5.12.0-py3-none-any.whl", hash = "sha256:7b1deeebbf351c7578e09bf2f63fa2ce8b5ffec296e0d349139d43cca061a81a"}, 376 | {file = "importlib_resources-5.12.0.tar.gz", hash = "sha256:4be82589bf5c1d7999aedf2a45159d10cb3ca4f19b2271f8792bc8e6da7b22f6"}, 377 | ] 378 | 379 | [package.dependencies] 380 | zipp = {version = ">=3.1.0", markers = "python_version < \"3.10\""} 381 | 382 | [package.extras] 383 | docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] 384 | testing = ["flake8 (<5)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)"] 385 | 386 | [[package]] 387 | name = "iniconfig" 388 | version = "2.0.0" 389 | description = "brain-dead simple config-ini parsing" 390 | category = "dev" 391 | optional = false 392 | python-versions = ">=3.7" 393 | files = [ 394 | {file = "iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374"}, 395 | {file = "iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3"}, 396 | ] 397 | 398 | [[package]] 399 | name = "isort" 400 | version = "5.12.0" 401 | description = "A Python utility / library to sort Python imports." 402 | category = "dev" 403 | optional = false 404 | python-versions = ">=3.8.0" 405 | files = [ 406 | {file = "isort-5.12.0-py3-none-any.whl", hash = "sha256:f84c2818376e66cf843d497486ea8fed8700b340f308f076c6fb1229dff318b6"}, 407 | {file = "isort-5.12.0.tar.gz", hash = "sha256:8bef7dde241278824a6d83f44a544709b065191b95b6e50894bdc722fcba0504"}, 408 | ] 409 | 410 | [package.extras] 411 | colors = ["colorama (>=0.4.3)"] 412 | pipfile-deprecated-finder = ["pip-shims (>=0.5.2)", "pipreqs", "requirementslib"] 413 | plugins = ["setuptools"] 414 | requirements-deprecated-finder = ["pip-api", "pipreqs"] 415 | 416 | [[package]] 417 | name = "jax" 418 | version = "0.4.8" 419 | description = "Differentiate, compile, and transform Numpy code." 420 | category = "main" 421 | optional = false 422 | python-versions = ">=3.8" 423 | files = [ 424 | {file = "jax-0.4.8.tar.gz", hash = "sha256:08116481f7336db16c24812bfb5e6f9786915f4c2f6ff4028331fa69e7535202"}, 425 | ] 426 | 427 | [package.dependencies] 428 | ml_dtypes = ">=0.0.3" 429 | numpy = ">=1.21" 430 | opt_einsum = "*" 431 | scipy = ">=1.7" 432 | 433 | [package.extras] 434 | australis = ["protobuf (>=3.13,<4)"] 435 | ci = ["jaxlib (==0.4.7)"] 436 | cpu = ["jaxlib (==0.4.7)"] 437 | cuda = ["jaxlib (==0.4.7+cuda11.cudnn86)"] 438 | cuda11-cudnn82 = ["jaxlib (==0.4.7+cuda11.cudnn82)"] 439 | cuda11-cudnn86 = ["jaxlib (==0.4.7+cuda11.cudnn86)"] 440 | cuda11-local = ["jaxlib (==0.4.7+cuda11.cudnn86)"] 441 | cuda11-pip = ["jaxlib (==0.4.7+cuda11.cudnn86)", "nvidia-cublas-cu11 (>=11.11)", "nvidia-cuda-nvcc-cu11 (>=11.8)", "nvidia-cuda-runtime-cu11 (>=11.8)", "nvidia-cudnn-cu11 (>=8.6)", "nvidia-cufft-cu11 (>=10.9)", "nvidia-cusolver-cu11 (>=11.4)", "nvidia-cusparse-cu11 (>=11.7)"] 442 | cuda12-local = ["jaxlib (==0.4.7+cuda12.cudnn88)"] 443 | cuda12-pip = ["jaxlib (==0.4.7+cuda12.cudnn88)", "nvidia-cublas-cu12", "nvidia-cuda-nvcc-cu12", "nvidia-cuda-runtime-cu12", "nvidia-cudnn-cu12", "nvidia-cufft-cu12", "nvidia-cusolver-cu12", "nvidia-cusparse-cu12"] 444 | minimum-jaxlib = ["jaxlib (==0.4.7)"] 445 | tpu = ["jaxlib (==0.4.7)", "libtpu-nightly (==0.1.dev20230327)", "requests"] 446 | 447 | [[package]] 448 | name = "jaxlib" 449 | version = "0.4.7" 450 | description = "XLA library for JAX" 451 | category = "main" 452 | optional = false 453 | python-versions = ">=3.8" 454 | files = [ 455 | {file = "jaxlib-0.4.7-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:63c2890978e8646516db3d8a680b43d2bed8b63543a70556391f589a261bd85f"}, 456 | {file = "jaxlib-0.4.7-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0c16f922507277d5630e81d9c1a4974366a27aad5230d645d063bc2011564d01"}, 457 | {file = "jaxlib-0.4.7-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:da88382e6487805974cea6facc61ba92b5828a7a1f2dd80f762c487d873a2b47"}, 458 | {file = "jaxlib-0.4.7-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:022b216036c009989d4c0683538820c19247215bb99fdd35c7bf32838d596be6"}, 459 | {file = "jaxlib-0.4.7-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d0f1d3b6ef6c68013898cca958ab1507d6809b523275037efbdb9aaaaab158ba"}, 460 | {file = "jaxlib-0.4.7-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:0ae7178c33460822d9d8d03718cba395e02e6bac2402709c35826c94f0c9cc7b"}, 461 | {file = "jaxlib-0.4.7-cp38-cp38-macosx_10_14_x86_64.whl", hash = "sha256:ea07605e37d2b4e25f3c639e0d22ab4605fbc1a10ea918fd14ce09077bdaffb6"}, 462 | {file = "jaxlib-0.4.7-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:48b85d3c8923b1619ddf8cbf14c4e4daf6919796d8aa9d006ce2a085e8202930"}, 463 | {file = "jaxlib-0.4.7-cp38-cp38-manylinux2014_x86_64.whl", hash = "sha256:a860f2990c97bee5ffcdbb5111751591e5e7a66d5e32b4f6d9e6aa14ac82bf27"}, 464 | {file = "jaxlib-0.4.7-cp39-cp39-macosx_10_14_x86_64.whl", hash = "sha256:c78dc2b6fa1c92ead137a23d1bd3e10d04c58b268e77eca811502abac05b2b19"}, 465 | {file = "jaxlib-0.4.7-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f1f3726e374d0d6fcc14da540b71b758d37356c6726f0f4b48e2f5530a5f8769"}, 466 | {file = "jaxlib-0.4.7-cp39-cp39-manylinux2014_x86_64.whl", hash = "sha256:d4629205dbe342153941db5f69c4a1bfe35fd8d2947aebe34f4dff3771d3fff7"}, 467 | ] 468 | 469 | [package.dependencies] 470 | ml-dtypes = ">=0.0.3" 471 | numpy = ">=1.21" 472 | scipy = ">=1.7" 473 | 474 | [[package]] 475 | name = "markdown-it-py" 476 | version = "2.2.0" 477 | description = "Python port of markdown-it. Markdown parsing, done right!" 478 | category = "dev" 479 | optional = false 480 | python-versions = ">=3.7" 481 | files = [ 482 | {file = "markdown-it-py-2.2.0.tar.gz", hash = "sha256:7c9a5e412688bc771c67432cbfebcdd686c93ce6484913dccf06cb5a0bea35a1"}, 483 | {file = "markdown_it_py-2.2.0-py3-none-any.whl", hash = "sha256:5a35f8d1870171d9acc47b99612dc146129b631baf04970128b568f190d0cc30"}, 484 | ] 485 | 486 | [package.dependencies] 487 | mdurl = ">=0.1,<1.0" 488 | 489 | [package.extras] 490 | benchmarking = ["psutil", "pytest", "pytest-benchmark"] 491 | code-style = ["pre-commit (>=3.0,<4.0)"] 492 | compare = ["commonmark (>=0.9,<1.0)", "markdown (>=3.4,<4.0)", "mistletoe (>=1.0,<2.0)", "mistune (>=2.0,<3.0)", "panflute (>=2.3,<3.0)"] 493 | linkify = ["linkify-it-py (>=1,<3)"] 494 | plugins = ["mdit-py-plugins"] 495 | profiling = ["gprof2dot"] 496 | rtd = ["attrs", "myst-parser", "pyyaml", "sphinx", "sphinx-copybutton", "sphinx-design", "sphinx_book_theme"] 497 | testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions"] 498 | 499 | [[package]] 500 | name = "mdurl" 501 | version = "0.1.2" 502 | description = "Markdown URL utilities" 503 | category = "dev" 504 | optional = false 505 | python-versions = ">=3.7" 506 | files = [ 507 | {file = "mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8"}, 508 | {file = "mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba"}, 509 | ] 510 | 511 | [[package]] 512 | name = "ml-dtypes" 513 | version = "0.0.4" 514 | description = "" 515 | category = "main" 516 | optional = false 517 | python-versions = ">=3.7" 518 | files = [ 519 | {file = "ml_dtypes-0.0.4-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a03c5acc55a878fac190d428ef01438f930cbef3fb8625c8c8fd2e3adc277607"}, 520 | {file = "ml_dtypes-0.0.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8e600aa70a9f8ee85c9488eb14852124c878ec824c3c7996d2d82010655eabfe"}, 521 | {file = "ml_dtypes-0.0.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a74c1fb29d2e586f643fb1a70b1dffe9fc35bc3ad8c76ec0797b2bf9f7ac128b"}, 522 | {file = "ml_dtypes-0.0.4-cp310-cp310-win_amd64.whl", hash = "sha256:b3f49901eb42cac259156edc17d4c1922ac47ddd1fe3c05169f445135a07319c"}, 523 | {file = "ml_dtypes-0.0.4-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:52aaa9318e2a4ec65a6bc4842df3442a9cfa00a9b8365a08e0370b0dfefc3a5a"}, 524 | {file = "ml_dtypes-0.0.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:db9912d50466d386a4016b16f889722183f6d6c03d9e478fdf62f41e50de0059"}, 525 | {file = "ml_dtypes-0.0.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4ece1269b5311489e26b3f3181d498b8829042f380cd160d7fe02f2393f69a71"}, 526 | {file = "ml_dtypes-0.0.4-cp311-cp311-win_amd64.whl", hash = "sha256:68d2e6c83c762aa6d476ea715ce6b2ac67f519c242cfe93d7a49cb76a83f6650"}, 527 | {file = "ml_dtypes-0.0.4-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:85085f9dac85b1eee5f7d2044c47bb3df72abc7785d38d176744fde5782b76ce"}, 528 | {file = "ml_dtypes-0.0.4-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a75ef23de72daf5efcc99799dfaa387386b79502a123909b0d3098ef84ffa6fa"}, 529 | {file = "ml_dtypes-0.0.4-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b651fa1f91ce83cf037db202cd2601ac9b649016ec8593459c0295e613bf47"}, 530 | {file = "ml_dtypes-0.0.4-cp38-cp38-win_amd64.whl", hash = "sha256:b28c6b7831fa2cbb3169ed3053f10fb11d0415e2f250b893eb874e3af747a1f3"}, 531 | {file = "ml_dtypes-0.0.4-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:23ff15cd9ba61cc42287097c30ae6841facd6dc14cc252f977d6430b8cd6eccc"}, 532 | {file = "ml_dtypes-0.0.4-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e5b148131da64f85053b79380cf34471eb869f7c027e2198a0c86d5e6fc9531f"}, 533 | {file = "ml_dtypes-0.0.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ebc64866c1848999fab6f4a2938e769aed95b964085ebdcd7cd45e350192e457"}, 534 | {file = "ml_dtypes-0.0.4-cp39-cp39-win_amd64.whl", hash = "sha256:e64869be11c830736c40513c47918c421a8385243846f1e8fd838793d866aa87"}, 535 | {file = "ml_dtypes-0.0.4.tar.gz", hash = "sha256:45623c738d477d7a0f3f8e4c94998dc49025202c520e62e27f0ef688db2f696f"}, 536 | ] 537 | 538 | [package.dependencies] 539 | numpy = [ 540 | {version = ">1.20", markers = "python_version <= \"3.9\""}, 541 | {version = ">=1.23.3", markers = "python_version > \"3.10\""}, 542 | {version = ">=1.21.2", markers = "python_version > \"3.9\""}, 543 | ] 544 | 545 | [package.extras] 546 | dev = ["absl-py", "pyink", "pylint (>=2.6.0)", "pytest", "pytest-xdist"] 547 | 548 | [[package]] 549 | name = "msgpack" 550 | version = "1.0.5" 551 | description = "MessagePack serializer" 552 | category = "dev" 553 | optional = false 554 | python-versions = "*" 555 | files = [ 556 | {file = "msgpack-1.0.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:525228efd79bb831cf6830a732e2e80bc1b05436b086d4264814b4b2955b2fa9"}, 557 | {file = "msgpack-1.0.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:4f8d8b3bf1ff2672567d6b5c725a1b347fe838b912772aa8ae2bf70338d5a198"}, 558 | {file = "msgpack-1.0.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:cdc793c50be3f01106245a61b739328f7dccc2c648b501e237f0699fe1395b81"}, 559 | {file = "msgpack-1.0.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5cb47c21a8a65b165ce29f2bec852790cbc04936f502966768e4aae9fa763cb7"}, 560 | {file = "msgpack-1.0.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e42b9594cc3bf4d838d67d6ed62b9e59e201862a25e9a157019e171fbe672dd3"}, 561 | {file = "msgpack-1.0.5-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:55b56a24893105dc52c1253649b60f475f36b3aa0fc66115bffafb624d7cb30b"}, 562 | {file = "msgpack-1.0.5-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:1967f6129fc50a43bfe0951c35acbb729be89a55d849fab7686004da85103f1c"}, 563 | {file = "msgpack-1.0.5-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:20a97bf595a232c3ee6d57ddaadd5453d174a52594bf9c21d10407e2a2d9b3bd"}, 564 | {file = "msgpack-1.0.5-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:d25dd59bbbbb996eacf7be6b4ad082ed7eacc4e8f3d2df1ba43822da9bfa122a"}, 565 | {file = "msgpack-1.0.5-cp310-cp310-win32.whl", hash = "sha256:382b2c77589331f2cb80b67cc058c00f225e19827dbc818d700f61513ab47bea"}, 566 | {file = "msgpack-1.0.5-cp310-cp310-win_amd64.whl", hash = "sha256:4867aa2df9e2a5fa5f76d7d5565d25ec76e84c106b55509e78c1ede0f152659a"}, 567 | {file = "msgpack-1.0.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:9f5ae84c5c8a857ec44dc180a8b0cc08238e021f57abdf51a8182e915e6299f0"}, 568 | {file = "msgpack-1.0.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:9e6ca5d5699bcd89ae605c150aee83b5321f2115695e741b99618f4856c50898"}, 569 | {file = "msgpack-1.0.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5494ea30d517a3576749cad32fa27f7585c65f5f38309c88c6d137877fa28a5a"}, 570 | {file = "msgpack-1.0.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1ab2f3331cb1b54165976a9d976cb251a83183631c88076613c6c780f0d6e45a"}, 571 | {file = "msgpack-1.0.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:28592e20bbb1620848256ebc105fc420436af59515793ed27d5c77a217477705"}, 572 | {file = "msgpack-1.0.5-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:fe5c63197c55bce6385d9aee16c4d0641684628f63ace85f73571e65ad1c1e8d"}, 573 | {file = "msgpack-1.0.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:ed40e926fa2f297e8a653c954b732f125ef97bdd4c889f243182299de27e2aa9"}, 574 | {file = "msgpack-1.0.5-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:b2de4c1c0538dcb7010902a2b97f4e00fc4ddf2c8cda9749af0e594d3b7fa3d7"}, 575 | {file = "msgpack-1.0.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:bf22a83f973b50f9d38e55c6aade04c41ddda19b00c4ebc558930d78eecc64ed"}, 576 | {file = "msgpack-1.0.5-cp311-cp311-win32.whl", hash = "sha256:c396e2cc213d12ce017b686e0f53497f94f8ba2b24799c25d913d46c08ec422c"}, 577 | {file = "msgpack-1.0.5-cp311-cp311-win_amd64.whl", hash = "sha256:6c4c68d87497f66f96d50142a2b73b97972130d93677ce930718f68828b382e2"}, 578 | {file = "msgpack-1.0.5-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:a2b031c2e9b9af485d5e3c4520f4220d74f4d222a5b8dc8c1a3ab9448ca79c57"}, 579 | {file = "msgpack-1.0.5-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4f837b93669ce4336e24d08286c38761132bc7ab29782727f8557e1eb21b2080"}, 580 | {file = "msgpack-1.0.5-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b1d46dfe3832660f53b13b925d4e0fa1432b00f5f7210eb3ad3bb9a13c6204a6"}, 581 | {file = "msgpack-1.0.5-cp36-cp36m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:366c9a7b9057e1547f4ad51d8facad8b406bab69c7d72c0eb6f529cf76d4b85f"}, 582 | {file = "msgpack-1.0.5-cp36-cp36m-musllinux_1_1_aarch64.whl", hash = "sha256:4c075728a1095efd0634a7dccb06204919a2f67d1893b6aa8e00497258bf926c"}, 583 | {file = "msgpack-1.0.5-cp36-cp36m-musllinux_1_1_i686.whl", hash = "sha256:f933bbda5a3ee63b8834179096923b094b76f0c7a73c1cfe8f07ad608c58844b"}, 584 | {file = "msgpack-1.0.5-cp36-cp36m-musllinux_1_1_x86_64.whl", hash = "sha256:36961b0568c36027c76e2ae3ca1132e35123dcec0706c4b7992683cc26c1320c"}, 585 | {file = "msgpack-1.0.5-cp36-cp36m-win32.whl", hash = "sha256:b5ef2f015b95f912c2fcab19c36814963b5463f1fb9049846994b007962743e9"}, 586 | {file = "msgpack-1.0.5-cp36-cp36m-win_amd64.whl", hash = "sha256:288e32b47e67f7b171f86b030e527e302c91bd3f40fd9033483f2cacc37f327a"}, 587 | {file = "msgpack-1.0.5-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:137850656634abddfb88236008339fdaba3178f4751b28f270d2ebe77a563b6c"}, 588 | {file = "msgpack-1.0.5-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0c05a4a96585525916b109bb85f8cb6511db1c6f5b9d9cbcbc940dc6b4be944b"}, 589 | {file = "msgpack-1.0.5-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:56a62ec00b636583e5cb6ad313bbed36bb7ead5fa3a3e38938503142c72cba4f"}, 590 | {file = "msgpack-1.0.5-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ef8108f8dedf204bb7b42994abf93882da1159728a2d4c5e82012edd92c9da9f"}, 591 | {file = "msgpack-1.0.5-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:1835c84d65f46900920b3708f5ba829fb19b1096c1800ad60bae8418652a951d"}, 592 | {file = "msgpack-1.0.5-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:e57916ef1bd0fee4f21c4600e9d1da352d8816b52a599c46460e93a6e9f17086"}, 593 | {file = "msgpack-1.0.5-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:17358523b85973e5f242ad74aa4712b7ee560715562554aa2134d96e7aa4cbbf"}, 594 | {file = "msgpack-1.0.5-cp37-cp37m-win32.whl", hash = "sha256:cb5aaa8c17760909ec6cb15e744c3ebc2ca8918e727216e79607b7bbce9c8f77"}, 595 | {file = "msgpack-1.0.5-cp37-cp37m-win_amd64.whl", hash = "sha256:ab31e908d8424d55601ad7075e471b7d0140d4d3dd3272daf39c5c19d936bd82"}, 596 | {file = "msgpack-1.0.5-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:b72d0698f86e8d9ddf9442bdedec15b71df3598199ba33322d9711a19f08145c"}, 597 | {file = "msgpack-1.0.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:379026812e49258016dd84ad79ac8446922234d498058ae1d415f04b522d5b2d"}, 598 | {file = "msgpack-1.0.5-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:332360ff25469c346a1c5e47cbe2a725517919892eda5cfaffe6046656f0b7bb"}, 599 | {file = "msgpack-1.0.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:476a8fe8fae289fdf273d6d2a6cb6e35b5a58541693e8f9f019bfe990a51e4ba"}, 600 | {file = "msgpack-1.0.5-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a9985b214f33311df47e274eb788a5893a761d025e2b92c723ba4c63936b69b1"}, 601 | {file = "msgpack-1.0.5-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:48296af57cdb1d885843afd73c4656be5c76c0c6328db3440c9601a98f303d87"}, 602 | {file = "msgpack-1.0.5-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:addab7e2e1fcc04bd08e4eb631c2a90960c340e40dfc4a5e24d2ff0d5a3b3edb"}, 603 | {file = "msgpack-1.0.5-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:916723458c25dfb77ff07f4c66aed34e47503b2eb3188b3adbec8d8aa6e00f48"}, 604 | {file = "msgpack-1.0.5-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:821c7e677cc6acf0fd3f7ac664c98803827ae6de594a9f99563e48c5a2f27eb0"}, 605 | {file = "msgpack-1.0.5-cp38-cp38-win32.whl", hash = "sha256:1c0f7c47f0087ffda62961d425e4407961a7ffd2aa004c81b9c07d9269512f6e"}, 606 | {file = "msgpack-1.0.5-cp38-cp38-win_amd64.whl", hash = "sha256:bae7de2026cbfe3782c8b78b0db9cbfc5455e079f1937cb0ab8d133496ac55e1"}, 607 | {file = "msgpack-1.0.5-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:20c784e66b613c7f16f632e7b5e8a1651aa5702463d61394671ba07b2fc9e025"}, 608 | {file = "msgpack-1.0.5-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:266fa4202c0eb94d26822d9bfd7af25d1e2c088927fe8de9033d929dd5ba24c5"}, 609 | {file = "msgpack-1.0.5-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:18334484eafc2b1aa47a6d42427da7fa8f2ab3d60b674120bce7a895a0a85bdd"}, 610 | {file = "msgpack-1.0.5-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:57e1f3528bd95cc44684beda696f74d3aaa8a5e58c816214b9046512240ef437"}, 611 | {file = "msgpack-1.0.5-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:586d0d636f9a628ddc6a17bfd45aa5b5efaf1606d2b60fa5d87b8986326e933f"}, 612 | {file = "msgpack-1.0.5-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a740fa0e4087a734455f0fc3abf5e746004c9da72fbd541e9b113013c8dc3282"}, 613 | {file = "msgpack-1.0.5-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:3055b0455e45810820db1f29d900bf39466df96ddca11dfa6d074fa47054376d"}, 614 | {file = "msgpack-1.0.5-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:a61215eac016f391129a013c9e46f3ab308db5f5ec9f25811e811f96962599a8"}, 615 | {file = "msgpack-1.0.5-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:362d9655cd369b08fda06b6657a303eb7172d5279997abe094512e919cf74b11"}, 616 | {file = "msgpack-1.0.5-cp39-cp39-win32.whl", hash = "sha256:ac9dd47af78cae935901a9a500104e2dea2e253207c924cc95de149606dc43cc"}, 617 | {file = "msgpack-1.0.5-cp39-cp39-win_amd64.whl", hash = "sha256:06f5174b5f8ed0ed919da0e62cbd4ffde676a374aba4020034da05fab67b9164"}, 618 | {file = "msgpack-1.0.5.tar.gz", hash = "sha256:c075544284eadc5cddc70f4757331d99dcbc16b2bbd4849d15f8aae4cf36d31c"}, 619 | ] 620 | 621 | [[package]] 622 | name = "mypy-extensions" 623 | version = "1.0.0" 624 | description = "Type system extensions for programs checked with the mypy type checker." 625 | category = "dev" 626 | optional = false 627 | python-versions = ">=3.5" 628 | files = [ 629 | {file = "mypy_extensions-1.0.0-py3-none-any.whl", hash = "sha256:4392f6c0eb8a5668a69e23d168ffa70f0be9ccfd32b5cc2d26a34ae5b844552d"}, 630 | {file = "mypy_extensions-1.0.0.tar.gz", hash = "sha256:75dbf8955dc00442a438fc4d0666508a9a97b6bd41aa2f0ffe9d2f2725af0782"}, 631 | ] 632 | 633 | [[package]] 634 | name = "nest-asyncio" 635 | version = "1.5.6" 636 | description = "Patch asyncio to allow nested event loops" 637 | category = "dev" 638 | optional = false 639 | python-versions = ">=3.5" 640 | files = [ 641 | {file = "nest_asyncio-1.5.6-py3-none-any.whl", hash = "sha256:b9a953fb40dceaa587d109609098db21900182b16440652454a146cffb06e8b8"}, 642 | {file = "nest_asyncio-1.5.6.tar.gz", hash = "sha256:d267cc1ff794403f7df692964d1d2a3fa9418ffea2a3f6859a439ff482fef290"}, 643 | ] 644 | 645 | [[package]] 646 | name = "nodeenv" 647 | version = "1.7.0" 648 | description = "Node.js virtual environment builder" 649 | category = "dev" 650 | optional = false 651 | python-versions = ">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*" 652 | files = [ 653 | {file = "nodeenv-1.7.0-py2.py3-none-any.whl", hash = "sha256:27083a7b96a25f2f5e1d8cb4b6317ee8aeda3bdd121394e5ac54e498028a042e"}, 654 | {file = "nodeenv-1.7.0.tar.gz", hash = "sha256:e0e7f7dfb85fc5394c6fe1e8fa98131a2473e04311a45afb6508f7cf1836fa2b"}, 655 | ] 656 | 657 | [package.dependencies] 658 | setuptools = "*" 659 | 660 | [[package]] 661 | name = "numpy" 662 | version = "1.24.2" 663 | description = "Fundamental package for array computing in Python" 664 | category = "main" 665 | optional = false 666 | python-versions = ">=3.8" 667 | files = [ 668 | {file = "numpy-1.24.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:eef70b4fc1e872ebddc38cddacc87c19a3709c0e3e5d20bf3954c147b1dd941d"}, 669 | {file = "numpy-1.24.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e8d2859428712785e8a8b7d2b3ef0a1d1565892367b32f915c4a4df44d0e64f5"}, 670 | {file = "numpy-1.24.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6524630f71631be2dabe0c541e7675db82651eb998496bbe16bc4f77f0772253"}, 671 | {file = "numpy-1.24.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a51725a815a6188c662fb66fb32077709a9ca38053f0274640293a14fdd22978"}, 672 | {file = "numpy-1.24.2-cp310-cp310-win32.whl", hash = "sha256:2620e8592136e073bd12ee4536149380695fbe9ebeae845b81237f986479ffc9"}, 673 | {file = "numpy-1.24.2-cp310-cp310-win_amd64.whl", hash = "sha256:97cf27e51fa078078c649a51d7ade3c92d9e709ba2bfb97493007103c741f1d0"}, 674 | {file = "numpy-1.24.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:7de8fdde0003f4294655aa5d5f0a89c26b9f22c0a58790c38fae1ed392d44a5a"}, 675 | {file = "numpy-1.24.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4173bde9fa2a005c2c6e2ea8ac1618e2ed2c1c6ec8a7657237854d42094123a0"}, 676 | {file = "numpy-1.24.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4cecaed30dc14123020f77b03601559fff3e6cd0c048f8b5289f4eeabb0eb281"}, 677 | {file = "numpy-1.24.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9a23f8440561a633204a67fb44617ce2a299beecf3295f0d13c495518908e910"}, 678 | {file = "numpy-1.24.2-cp311-cp311-win32.whl", hash = "sha256:e428c4fbfa085f947b536706a2fc349245d7baa8334f0c5723c56a10595f9b95"}, 679 | {file = "numpy-1.24.2-cp311-cp311-win_amd64.whl", hash = "sha256:557d42778a6869c2162deb40ad82612645e21d79e11c1dc62c6e82a2220ffb04"}, 680 | {file = "numpy-1.24.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:d0a2db9d20117bf523dde15858398e7c0858aadca7c0f088ac0d6edd360e9ad2"}, 681 | {file = "numpy-1.24.2-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:c72a6b2f4af1adfe193f7beb91ddf708ff867a3f977ef2ec53c0ffb8283ab9f5"}, 682 | {file = "numpy-1.24.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c29e6bd0ec49a44d7690ecb623a8eac5ab8a923bce0bea6293953992edf3a76a"}, 683 | {file = "numpy-1.24.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2eabd64ddb96a1239791da78fa5f4e1693ae2dadc82a76bc76a14cbb2b966e96"}, 684 | {file = "numpy-1.24.2-cp38-cp38-win32.whl", hash = "sha256:e3ab5d32784e843fc0dd3ab6dcafc67ef806e6b6828dc6af2f689be0eb4d781d"}, 685 | {file = "numpy-1.24.2-cp38-cp38-win_amd64.whl", hash = "sha256:76807b4063f0002c8532cfeac47a3068a69561e9c8715efdad3c642eb27c0756"}, 686 | {file = "numpy-1.24.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:4199e7cfc307a778f72d293372736223e39ec9ac096ff0a2e64853b866a8e18a"}, 687 | {file = "numpy-1.24.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:adbdce121896fd3a17a77ab0b0b5eedf05a9834a18699db6829a64e1dfccca7f"}, 688 | {file = "numpy-1.24.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:889b2cc88b837d86eda1b17008ebeb679d82875022200c6e8e4ce6cf549b7acb"}, 689 | {file = "numpy-1.24.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f64bb98ac59b3ea3bf74b02f13836eb2e24e48e0ab0145bbda646295769bd780"}, 690 | {file = "numpy-1.24.2-cp39-cp39-win32.whl", hash = "sha256:63e45511ee4d9d976637d11e6c9864eae50e12dc9598f531c035265991910468"}, 691 | {file = "numpy-1.24.2-cp39-cp39-win_amd64.whl", hash = "sha256:a77d3e1163a7770164404607b7ba3967fb49b24782a6ef85d9b5f54126cc39e5"}, 692 | {file = "numpy-1.24.2-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:92011118955724465fb6853def593cf397b4a1367495e0b59a7e69d40c4eb71d"}, 693 | {file = "numpy-1.24.2-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f9006288bcf4895917d02583cf3411f98631275bc67cce355a7f39f8c14338fa"}, 694 | {file = "numpy-1.24.2-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:150947adbdfeceec4e5926d956a06865c1c690f2fd902efede4ca6fe2e657c3f"}, 695 | {file = "numpy-1.24.2.tar.gz", hash = "sha256:003a9f530e880cb2cd177cba1af7220b9aa42def9c4afc2a2fc3ee6be7eb2b22"}, 696 | ] 697 | 698 | [[package]] 699 | name = "opt-einsum" 700 | version = "3.3.0" 701 | description = "Optimizing numpys einsum function" 702 | category = "main" 703 | optional = false 704 | python-versions = ">=3.5" 705 | files = [ 706 | {file = "opt_einsum-3.3.0-py3-none-any.whl", hash = "sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147"}, 707 | {file = "opt_einsum-3.3.0.tar.gz", hash = "sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549"}, 708 | ] 709 | 710 | [package.dependencies] 711 | numpy = ">=1.7" 712 | 713 | [package.extras] 714 | docs = ["numpydoc", "sphinx (==1.2.3)", "sphinx-rtd-theme", "sphinxcontrib-napoleon"] 715 | tests = ["pytest", "pytest-cov", "pytest-pep8"] 716 | 717 | [[package]] 718 | name = "optax" 719 | version = "0.1.4" 720 | description = "A gradient processing and optimisation library in JAX." 721 | category = "dev" 722 | optional = false 723 | python-versions = ">=3.7" 724 | files = [ 725 | {file = "optax-0.1.4-py3-none-any.whl", hash = "sha256:12fcf33bd682f9a162a3deb097f864130c3224d76771af2ba09410de80399a9b"}, 726 | {file = "optax-0.1.4.tar.gz", hash = "sha256:fb7a0550d57a6636164a3de25986a8a19be8ff6431fcdf1225b4e05175810f22"}, 727 | ] 728 | 729 | [package.dependencies] 730 | absl-py = ">=0.7.1" 731 | chex = ">=0.1.5" 732 | jax = ">=0.1.55" 733 | jaxlib = ">=0.1.37" 734 | numpy = ">=1.18.0" 735 | typing-extensions = ">=3.10.0" 736 | 737 | [[package]] 738 | name = "orbax" 739 | version = "0.1.7" 740 | description = "Orbax" 741 | category = "dev" 742 | optional = false 743 | python-versions = ">=3.8" 744 | files = [ 745 | {file = "orbax-0.1.7-py3-none-any.whl", hash = "sha256:67c7ce52b5476202af84977e8db03dede6c009b5d1f1095acfc175578038449b"}, 746 | {file = "orbax-0.1.7.tar.gz", hash = "sha256:2517f566134db6597d2850450b7f486efd24bf24962bc4881007f4cc8e978b37"}, 747 | ] 748 | 749 | [package.dependencies] 750 | absl-py = "*" 751 | cached_property = "*" 752 | etils = "*" 753 | importlib_resources = "*" 754 | jax = ">=0.4.6" 755 | jaxlib = "*" 756 | msgpack = "*" 757 | nest_asyncio = "*" 758 | numpy = "*" 759 | pyyaml = "*" 760 | tensorstore = ">=0.1.20" 761 | typing_extensions = "*" 762 | 763 | [package.extras] 764 | dev = ["flax", "pytest", "pytest-xdist"] 765 | 766 | [[package]] 767 | name = "packaging" 768 | version = "23.0" 769 | description = "Core utilities for Python packages" 770 | category = "dev" 771 | optional = false 772 | python-versions = ">=3.7" 773 | files = [ 774 | {file = "packaging-23.0-py3-none-any.whl", hash = "sha256:714ac14496c3e68c99c29b00845f7a2b85f3bb6f1078fd9f72fd20f0570002b2"}, 775 | {file = "packaging-23.0.tar.gz", hash = "sha256:b6ad297f8907de0fa2fe1ccbd26fdaf387f5f47c7275fedf8cce89f99446cf97"}, 776 | ] 777 | 778 | [[package]] 779 | name = "pathspec" 780 | version = "0.11.1" 781 | description = "Utility library for gitignore style pattern matching of file paths." 782 | category = "dev" 783 | optional = false 784 | python-versions = ">=3.7" 785 | files = [ 786 | {file = "pathspec-0.11.1-py3-none-any.whl", hash = "sha256:d8af70af76652554bd134c22b3e8a1cc46ed7d91edcdd721ef1a0c51a84a5293"}, 787 | {file = "pathspec-0.11.1.tar.gz", hash = "sha256:2798de800fa92780e33acca925945e9a19a133b715067cf165b8866c15a31687"}, 788 | ] 789 | 790 | [[package]] 791 | name = "platformdirs" 792 | version = "3.2.0" 793 | description = "A small Python package for determining appropriate platform-specific dirs, e.g. a \"user data dir\"." 794 | category = "dev" 795 | optional = false 796 | python-versions = ">=3.7" 797 | files = [ 798 | {file = "platformdirs-3.2.0-py3-none-any.whl", hash = "sha256:ebe11c0d7a805086e99506aa331612429a72ca7cd52a1f0d277dc4adc20cb10e"}, 799 | {file = "platformdirs-3.2.0.tar.gz", hash = "sha256:d5b638ca397f25f979350ff789db335903d7ea010ab28903f57b27e1b16c2b08"}, 800 | ] 801 | 802 | [package.extras] 803 | docs = ["furo (>=2022.12.7)", "proselint (>=0.13)", "sphinx (>=6.1.3)", "sphinx-autodoc-typehints (>=1.22,!=1.23.4)"] 804 | test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=7.2.2)", "pytest-cov (>=4)", "pytest-mock (>=3.10)"] 805 | 806 | [[package]] 807 | name = "pluggy" 808 | version = "1.0.0" 809 | description = "plugin and hook calling mechanisms for python" 810 | category = "dev" 811 | optional = false 812 | python-versions = ">=3.6" 813 | files = [ 814 | {file = "pluggy-1.0.0-py2.py3-none-any.whl", hash = "sha256:74134bbf457f031a36d68416e1509f34bd5ccc019f0bcc952c7b909d06b37bd3"}, 815 | {file = "pluggy-1.0.0.tar.gz", hash = "sha256:4224373bacce55f955a878bf9cfa763c1e360858e330072059e10bad68531159"}, 816 | ] 817 | 818 | [package.extras] 819 | dev = ["pre-commit", "tox"] 820 | testing = ["pytest", "pytest-benchmark"] 821 | 822 | [[package]] 823 | name = "pre-commit" 824 | version = "3.2.2" 825 | description = "A framework for managing and maintaining multi-language pre-commit hooks." 826 | category = "dev" 827 | optional = false 828 | python-versions = ">=3.8" 829 | files = [ 830 | {file = "pre_commit-3.2.2-py2.py3-none-any.whl", hash = "sha256:0b4210aea813fe81144e87c5a291f09ea66f199f367fa1df41b55e1d26e1e2b4"}, 831 | {file = "pre_commit-3.2.2.tar.gz", hash = "sha256:5b808fcbda4afbccf6d6633a56663fed35b6c2bc08096fd3d47ce197ac351d9d"}, 832 | ] 833 | 834 | [package.dependencies] 835 | cfgv = ">=2.0.0" 836 | identify = ">=1.0.0" 837 | nodeenv = ">=0.11.1" 838 | pyyaml = ">=5.1" 839 | virtualenv = ">=20.10.0" 840 | 841 | [[package]] 842 | name = "pygments" 843 | version = "2.15.0" 844 | description = "Pygments is a syntax highlighting package written in Python." 845 | category = "dev" 846 | optional = false 847 | python-versions = ">=3.7" 848 | files = [ 849 | {file = "Pygments-2.15.0-py3-none-any.whl", hash = "sha256:77a3299119af881904cd5ecd1ac6a66214b6e9bed1f2db16993b54adede64094"}, 850 | {file = "Pygments-2.15.0.tar.gz", hash = "sha256:f7e36cffc4c517fbc252861b9a6e4644ca0e5abadf9a113c72d1358ad09b9500"}, 851 | ] 852 | 853 | [package.extras] 854 | plugins = ["importlib-metadata"] 855 | 856 | [[package]] 857 | name = "pytest" 858 | version = "7.3.0" 859 | description = "pytest: simple powerful testing with Python" 860 | category = "dev" 861 | optional = false 862 | python-versions = ">=3.7" 863 | files = [ 864 | {file = "pytest-7.3.0-py3-none-any.whl", hash = "sha256:933051fa1bfbd38a21e73c3960cebdad4cf59483ddba7696c48509727e17f201"}, 865 | {file = "pytest-7.3.0.tar.gz", hash = "sha256:58ecc27ebf0ea643ebfdf7fb1249335da761a00c9f955bcd922349bcb68ee57d"}, 866 | ] 867 | 868 | [package.dependencies] 869 | colorama = {version = "*", markers = "sys_platform == \"win32\""} 870 | exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""} 871 | iniconfig = "*" 872 | packaging = "*" 873 | pluggy = ">=0.12,<2.0" 874 | tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""} 875 | 876 | [package.extras] 877 | testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "xmlschema"] 878 | 879 | [[package]] 880 | name = "pytest-cov" 881 | version = "4.0.0" 882 | description = "Pytest plugin for measuring coverage." 883 | category = "dev" 884 | optional = false 885 | python-versions = ">=3.6" 886 | files = [ 887 | {file = "pytest-cov-4.0.0.tar.gz", hash = "sha256:996b79efde6433cdbd0088872dbc5fb3ed7fe1578b68cdbba634f14bb8dd0470"}, 888 | {file = "pytest_cov-4.0.0-py3-none-any.whl", hash = "sha256:2feb1b751d66a8bd934e5edfa2e961d11309dc37b73b0eabe73b5945fee20f6b"}, 889 | ] 890 | 891 | [package.dependencies] 892 | coverage = {version = ">=5.2.1", extras = ["toml"]} 893 | pytest = ">=4.6" 894 | 895 | [package.extras] 896 | testing = ["fields", "hunter", "process-tests", "pytest-xdist", "six", "virtualenv"] 897 | 898 | [[package]] 899 | name = "pyyaml" 900 | version = "6.0" 901 | description = "YAML parser and emitter for Python" 902 | category = "dev" 903 | optional = false 904 | python-versions = ">=3.6" 905 | files = [ 906 | {file = "PyYAML-6.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d4db7c7aef085872ef65a8fd7d6d09a14ae91f691dec3e87ee5ee0539d516f53"}, 907 | {file = "PyYAML-6.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:9df7ed3b3d2e0ecfe09e14741b857df43adb5a3ddadc919a2d94fbdf78fea53c"}, 908 | {file = "PyYAML-6.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:77f396e6ef4c73fdc33a9157446466f1cff553d979bd00ecb64385760c6babdc"}, 909 | {file = "PyYAML-6.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a80a78046a72361de73f8f395f1f1e49f956c6be882eed58505a15f3e430962b"}, 910 | {file = "PyYAML-6.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:f84fbc98b019fef2ee9a1cb3ce93e3187a6df0b2538a651bfb890254ba9f90b5"}, 911 | {file = "PyYAML-6.0-cp310-cp310-win32.whl", hash = "sha256:2cd5df3de48857ed0544b34e2d40e9fac445930039f3cfe4bcc592a1f836d513"}, 912 | {file = "PyYAML-6.0-cp310-cp310-win_amd64.whl", hash = "sha256:daf496c58a8c52083df09b80c860005194014c3698698d1a57cbcfa182142a3a"}, 913 | {file = "PyYAML-6.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d4b0ba9512519522b118090257be113b9468d804b19d63c71dbcf4a48fa32358"}, 914 | {file = "PyYAML-6.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:81957921f441d50af23654aa6c5e5eaf9b06aba7f0a19c18a538dc7ef291c5a1"}, 915 | {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:afa17f5bc4d1b10afd4466fd3a44dc0e245382deca5b3c353d8b757f9e3ecb8d"}, 916 | {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dbad0e9d368bb989f4515da330b88a057617d16b6a8245084f1b05400f24609f"}, 917 | {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:432557aa2c09802be39460360ddffd48156e30721f5e8d917f01d31694216782"}, 918 | {file = "PyYAML-6.0-cp311-cp311-win32.whl", hash = "sha256:bfaef573a63ba8923503d27530362590ff4f576c626d86a9fed95822a8255fd7"}, 919 | {file = "PyYAML-6.0-cp311-cp311-win_amd64.whl", hash = "sha256:01b45c0191e6d66c470b6cf1b9531a771a83c1c4208272ead47a3ae4f2f603bf"}, 920 | {file = "PyYAML-6.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:897b80890765f037df3403d22bab41627ca8811ae55e9a722fd0392850ec4d86"}, 921 | {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:50602afada6d6cbfad699b0c7bb50d5ccffa7e46a3d738092afddc1f9758427f"}, 922 | {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:48c346915c114f5fdb3ead70312bd042a953a8ce5c7106d5bfb1a5254e47da92"}, 923 | {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:98c4d36e99714e55cfbaaee6dd5badbc9a1ec339ebfc3b1f52e293aee6bb71a4"}, 924 | {file = "PyYAML-6.0-cp36-cp36m-win32.whl", hash = "sha256:0283c35a6a9fbf047493e3a0ce8d79ef5030852c51e9d911a27badfde0605293"}, 925 | {file = "PyYAML-6.0-cp36-cp36m-win_amd64.whl", hash = "sha256:07751360502caac1c067a8132d150cf3d61339af5691fe9e87803040dbc5db57"}, 926 | {file = "PyYAML-6.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:819b3830a1543db06c4d4b865e70ded25be52a2e0631ccd2f6a47a2822f2fd7c"}, 927 | {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:473f9edb243cb1935ab5a084eb238d842fb8f404ed2193a915d1784b5a6b5fc0"}, 928 | {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0ce82d761c532fe4ec3f87fc45688bdd3a4c1dc5e0b4a19814b9009a29baefd4"}, 929 | {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:231710d57adfd809ef5d34183b8ed1eeae3f76459c18fb4a0b373ad56bedcdd9"}, 930 | {file = "PyYAML-6.0-cp37-cp37m-win32.whl", hash = "sha256:c5687b8d43cf58545ade1fe3e055f70eac7a5a1a0bf42824308d868289a95737"}, 931 | {file = "PyYAML-6.0-cp37-cp37m-win_amd64.whl", hash = "sha256:d15a181d1ecd0d4270dc32edb46f7cb7733c7c508857278d3d378d14d606db2d"}, 932 | {file = "PyYAML-6.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:0b4624f379dab24d3725ffde76559cff63d9ec94e1736b556dacdfebe5ab6d4b"}, 933 | {file = "PyYAML-6.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:213c60cd50106436cc818accf5baa1aba61c0189ff610f64f4a3e8c6726218ba"}, 934 | {file = "PyYAML-6.0-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9fa600030013c4de8165339db93d182b9431076eb98eb40ee068700c9c813e34"}, 935 | {file = "PyYAML-6.0-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:277a0ef2981ca40581a47093e9e2d13b3f1fbbeffae064c1d21bfceba2030287"}, 936 | {file = "PyYAML-6.0-cp38-cp38-win32.whl", hash = "sha256:d4eccecf9adf6fbcc6861a38015c2a64f38b9d94838ac1810a9023a0609e1b78"}, 937 | {file = "PyYAML-6.0-cp38-cp38-win_amd64.whl", hash = "sha256:1e4747bc279b4f613a09eb64bba2ba602d8a6664c6ce6396a4d0cd413a50ce07"}, 938 | {file = "PyYAML-6.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:055d937d65826939cb044fc8c9b08889e8c743fdc6a32b33e2390f66013e449b"}, 939 | {file = "PyYAML-6.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:e61ceaab6f49fb8bdfaa0f92c4b57bcfbea54c09277b1b4f7ac376bfb7a7c174"}, 940 | {file = "PyYAML-6.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d67d839ede4ed1b28a4e8909735fc992a923cdb84e618544973d7dfc71540803"}, 941 | {file = "PyYAML-6.0-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cba8c411ef271aa037d7357a2bc8f9ee8b58b9965831d9e51baf703280dc73d3"}, 942 | {file = "PyYAML-6.0-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:40527857252b61eacd1d9af500c3337ba8deb8fc298940291486c465c8b46ec0"}, 943 | {file = "PyYAML-6.0-cp39-cp39-win32.whl", hash = "sha256:b5b9eccad747aabaaffbc6064800670f0c297e52c12754eb1d976c57e4f74dcb"}, 944 | {file = "PyYAML-6.0-cp39-cp39-win_amd64.whl", hash = "sha256:b3d267842bf12586ba6c734f89d1f5b871df0273157918b0ccefa29deb05c21c"}, 945 | {file = "PyYAML-6.0.tar.gz", hash = "sha256:68fb519c14306fec9720a2a5b45bc9f0c8d1b9c72adf45c37baedfcd949c35a2"}, 946 | ] 947 | 948 | [[package]] 949 | name = "rich" 950 | version = "13.3.3" 951 | description = "Render rich text, tables, progress bars, syntax highlighting, markdown and more to the terminal" 952 | category = "dev" 953 | optional = false 954 | python-versions = ">=3.7.0" 955 | files = [ 956 | {file = "rich-13.3.3-py3-none-any.whl", hash = "sha256:540c7d6d26a1178e8e8b37e9ba44573a3cd1464ff6348b99ee7061b95d1c6333"}, 957 | {file = "rich-13.3.3.tar.gz", hash = "sha256:dc84400a9d842b3a9c5ff74addd8eb798d155f36c1c91303888e0a66850d2a15"}, 958 | ] 959 | 960 | [package.dependencies] 961 | markdown-it-py = ">=2.2.0,<3.0.0" 962 | pygments = ">=2.13.0,<3.0.0" 963 | typing-extensions = {version = ">=4.0.0,<5.0", markers = "python_version < \"3.9\""} 964 | 965 | [package.extras] 966 | jupyter = ["ipywidgets (>=7.5.1,<9)"] 967 | 968 | [[package]] 969 | name = "scipy" 970 | version = "1.10.1" 971 | description = "Fundamental algorithms for scientific computing in Python" 972 | category = "main" 973 | optional = false 974 | python-versions = "<3.12,>=3.8" 975 | files = [ 976 | {file = "scipy-1.10.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e7354fd7527a4b0377ce55f286805b34e8c54b91be865bac273f527e1b839019"}, 977 | {file = "scipy-1.10.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:4b3f429188c66603a1a5c549fb414e4d3bdc2a24792e061ffbd607d3d75fd84e"}, 978 | {file = "scipy-1.10.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1553b5dcddd64ba9a0d95355e63fe6c3fc303a8fd77c7bc91e77d61363f7433f"}, 979 | {file = "scipy-1.10.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4c0ff64b06b10e35215abce517252b375e580a6125fd5fdf6421b98efbefb2d2"}, 980 | {file = "scipy-1.10.1-cp310-cp310-win_amd64.whl", hash = "sha256:fae8a7b898c42dffe3f7361c40d5952b6bf32d10c4569098d276b4c547905ee1"}, 981 | {file = "scipy-1.10.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0f1564ea217e82c1bbe75ddf7285ba0709ecd503f048cb1236ae9995f64217bd"}, 982 | {file = "scipy-1.10.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:d925fa1c81b772882aa55bcc10bf88324dadb66ff85d548c71515f6689c6dac5"}, 983 | {file = "scipy-1.10.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aaea0a6be54462ec027de54fca511540980d1e9eea68b2d5c1dbfe084797be35"}, 984 | {file = "scipy-1.10.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15a35c4242ec5f292c3dd364a7c71a61be87a3d4ddcc693372813c0b73c9af1d"}, 985 | {file = "scipy-1.10.1-cp311-cp311-win_amd64.whl", hash = "sha256:43b8e0bcb877faf0abfb613d51026cd5cc78918e9530e375727bf0625c82788f"}, 986 | {file = "scipy-1.10.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5678f88c68ea866ed9ebe3a989091088553ba12c6090244fdae3e467b1139c35"}, 987 | {file = "scipy-1.10.1-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:39becb03541f9e58243f4197584286e339029e8908c46f7221abeea4b749fa88"}, 988 | {file = "scipy-1.10.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bce5869c8d68cf383ce240e44c1d9ae7c06078a9396df68ce88a1230f93a30c1"}, 989 | {file = "scipy-1.10.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:07c3457ce0b3ad5124f98a86533106b643dd811dd61b548e78cf4c8786652f6f"}, 990 | {file = "scipy-1.10.1-cp38-cp38-win_amd64.whl", hash = "sha256:049a8bbf0ad95277ffba9b3b7d23e5369cc39e66406d60422c8cfef40ccc8415"}, 991 | {file = "scipy-1.10.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:cd9f1027ff30d90618914a64ca9b1a77a431159df0e2a195d8a9e8a04c78abf9"}, 992 | {file = "scipy-1.10.1-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:79c8e5a6c6ffaf3a2262ef1be1e108a035cf4f05c14df56057b64acc5bebffb6"}, 993 | {file = "scipy-1.10.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:51af417a000d2dbe1ec6c372dfe688e041a7084da4fdd350aeb139bd3fb55353"}, 994 | {file = "scipy-1.10.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1b4735d6c28aad3cdcf52117e0e91d6b39acd4272f3f5cd9907c24ee931ad601"}, 995 | {file = "scipy-1.10.1-cp39-cp39-win_amd64.whl", hash = "sha256:7ff7f37b1bf4417baca958d254e8e2875d0cc23aaadbe65b3d5b3077b0eb23ea"}, 996 | {file = "scipy-1.10.1.tar.gz", hash = "sha256:2cf9dfb80a7b4589ba4c40ce7588986d6d5cebc5457cad2c2880f6bc2d42f3a5"}, 997 | ] 998 | 999 | [package.dependencies] 1000 | numpy = ">=1.19.5,<1.27.0" 1001 | 1002 | [package.extras] 1003 | dev = ["click", "doit (>=0.36.0)", "flake8", "mypy", "pycodestyle", "pydevtool", "rich-click", "typing_extensions"] 1004 | doc = ["matplotlib (>2)", "numpydoc", "pydata-sphinx-theme (==0.9.0)", "sphinx (!=4.1.0)", "sphinx-design (>=0.2.0)"] 1005 | test = ["asv", "gmpy2", "mpmath", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"] 1006 | 1007 | [[package]] 1008 | name = "setuptools" 1009 | version = "67.6.1" 1010 | description = "Easily download, build, install, upgrade, and uninstall Python packages" 1011 | category = "dev" 1012 | optional = false 1013 | python-versions = ">=3.7" 1014 | files = [ 1015 | {file = "setuptools-67.6.1-py3-none-any.whl", hash = "sha256:e728ca814a823bf7bf60162daf9db95b93d532948c4c0bea762ce62f60189078"}, 1016 | {file = "setuptools-67.6.1.tar.gz", hash = "sha256:257de92a9d50a60b8e22abfcbb771571fde0dbf3ec234463212027a4eeecbe9a"}, 1017 | ] 1018 | 1019 | [package.extras] 1020 | docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-hoverxref (<2)", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (==0.8.3)", "sphinx-reredirects", "sphinxcontrib-towncrier"] 1021 | testing = ["build[virtualenv]", "filelock (>=3.4.0)", "flake8 (<5)", "flake8-2020", "ini2toml[lite] (>=0.9)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pip (>=19.1)", "pip-run (>=8.8)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)", "pytest-perf", "pytest-timeout", "pytest-xdist", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"] 1022 | testing-integration = ["build[virtualenv]", "filelock (>=3.4.0)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pytest", "pytest-enabler", "pytest-xdist", "tomli", "virtualenv (>=13.0.0)", "wheel"] 1023 | 1024 | [[package]] 1025 | name = "tensorstore" 1026 | version = "0.1.35" 1027 | description = "Read and write large, multi-dimensional arrays" 1028 | category = "dev" 1029 | optional = false 1030 | python-versions = ">=3.8" 1031 | files = [ 1032 | {file = "tensorstore-0.1.35-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:4e1b4210b777c4a585a183bdd435bfa8aa5628c46075cb64adcd9b4bdd124e35"}, 1033 | {file = "tensorstore-0.1.35-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:285cfc4816bb0cc305cddc11f25f81a14bd84af0c8bbd39e42c81413c0bf242e"}, 1034 | {file = "tensorstore-0.1.35-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:51a8578518362daee85e39162bf870faf5545868cd53fe99ede0eda1fad288f6"}, 1035 | {file = "tensorstore-0.1.35-cp310-cp310-win_amd64.whl", hash = "sha256:a0318ea4afd4f2c00ce2dd4b540acb31c45a260bda94ae7e4340a1a4d28c6848"}, 1036 | {file = "tensorstore-0.1.35-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:7173f451c1b970230f57b0cdefdd692940ee457d4982a696d00aba163a7fee9a"}, 1037 | {file = "tensorstore-0.1.35-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d2fc9d8e5aaa54c538434592eaf88f7dfa6773fb35a960cc4cbe20bef55092d7"}, 1038 | {file = "tensorstore-0.1.35-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f34fc72a9ceff1c1666b6c85cc86858adbd10ff1c3c6b98c7788bbea9161a6a6"}, 1039 | {file = "tensorstore-0.1.35-cp311-cp311-win_amd64.whl", hash = "sha256:216bf4c00ec4aabf699d2a54ee9311f2fb19a2a3a904d7abb2194572af2f8384"}, 1040 | {file = "tensorstore-0.1.35-cp38-cp38-macosx_10_14_x86_64.whl", hash = "sha256:932f96d35ebdf0e4650bd9cd089319e9c5723d2aaf3f65123a821fc3b04ca4ac"}, 1041 | {file = "tensorstore-0.1.35-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:c9d0ea7888a88e5d892894ddf14f76d37665d0adbc2d861c559d0d6d5eaac20d"}, 1042 | {file = "tensorstore-0.1.35-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bd1406fd325517331887b0b2533e27e17e549196418eb087730bafa6fffaa236"}, 1043 | {file = "tensorstore-0.1.35-cp38-cp38-win_amd64.whl", hash = "sha256:99ad4e577249c2dfb07a501d78a31e29d6b9e53752384d58782ad53e6014ac41"}, 1044 | {file = "tensorstore-0.1.35-cp39-cp39-macosx_10_14_x86_64.whl", hash = "sha256:9acf9e9a7b3117881ec11f26930d0fee89cce6bb3d81056c15317f7cf2c0c1e1"}, 1045 | {file = "tensorstore-0.1.35-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:d742468c6aec6a1dcf3ec164694c9827eb498b2c701b8020fb4a56446f9bbc1a"}, 1046 | {file = "tensorstore-0.1.35-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f8011a556b4cbbc547e9365cf833ce7a5fb9fdbaa3f86a92665b78a78b78131"}, 1047 | {file = "tensorstore-0.1.35-cp39-cp39-win_amd64.whl", hash = "sha256:04c383af4a17e238fffdda7abc06b66c8583554d523cd721f2011bbb7a715327"}, 1048 | {file = "tensorstore-0.1.35.tar.gz", hash = "sha256:93db16e2f448cad716628640d3b73b87d9b259ae8ba1741a82108aef14e427c6"}, 1049 | ] 1050 | 1051 | [package.dependencies] 1052 | numpy = ">=1.16.0" 1053 | 1054 | [[package]] 1055 | name = "tomli" 1056 | version = "2.0.1" 1057 | description = "A lil' TOML parser" 1058 | category = "dev" 1059 | optional = false 1060 | python-versions = ">=3.7" 1061 | files = [ 1062 | {file = "tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc"}, 1063 | {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, 1064 | ] 1065 | 1066 | [[package]] 1067 | name = "toolz" 1068 | version = "0.12.0" 1069 | description = "List processing tools and functional utilities" 1070 | category = "dev" 1071 | optional = false 1072 | python-versions = ">=3.5" 1073 | files = [ 1074 | {file = "toolz-0.12.0-py3-none-any.whl", hash = "sha256:2059bd4148deb1884bb0eb770a3cde70e7f954cfbbdc2285f1f2de01fd21eb6f"}, 1075 | {file = "toolz-0.12.0.tar.gz", hash = "sha256:88c570861c440ee3f2f6037c4654613228ff40c93a6c25e0eba70d17282c6194"}, 1076 | ] 1077 | 1078 | [[package]] 1079 | name = "typing-extensions" 1080 | version = "4.5.0" 1081 | description = "Backported and Experimental Type Hints for Python 3.7+" 1082 | category = "main" 1083 | optional = false 1084 | python-versions = ">=3.7" 1085 | files = [ 1086 | {file = "typing_extensions-4.5.0-py3-none-any.whl", hash = "sha256:fb33085c39dd998ac16d1431ebc293a8b3eedd00fd4a32de0ff79002c19511b4"}, 1087 | {file = "typing_extensions-4.5.0.tar.gz", hash = "sha256:5cb5f4a79139d699607b3ef622a1dedafa84e115ab0024e0d9c044a9479ca7cb"}, 1088 | ] 1089 | 1090 | [[package]] 1091 | name = "virtualenv" 1092 | version = "20.21.0" 1093 | description = "Virtual Python Environment builder" 1094 | category = "dev" 1095 | optional = false 1096 | python-versions = ">=3.7" 1097 | files = [ 1098 | {file = "virtualenv-20.21.0-py3-none-any.whl", hash = "sha256:31712f8f2a17bd06234fa97fdf19609e789dd4e3e4bf108c3da71d710651adbc"}, 1099 | {file = "virtualenv-20.21.0.tar.gz", hash = "sha256:f50e3e60f990a0757c9b68333c9fdaa72d7188caa417f96af9e52407831a3b68"}, 1100 | ] 1101 | 1102 | [package.dependencies] 1103 | distlib = ">=0.3.6,<1" 1104 | filelock = ">=3.4.1,<4" 1105 | platformdirs = ">=2.4,<4" 1106 | 1107 | [package.extras] 1108 | docs = ["furo (>=2022.12.7)", "proselint (>=0.13)", "sphinx (>=6.1.3)", "sphinx-argparse (>=0.4)", "sphinxcontrib-towncrier (>=0.2.1a0)", "towncrier (>=22.12)"] 1109 | test = ["covdefaults (>=2.2.2)", "coverage (>=7.1)", "coverage-enable-subprocess (>=1)", "flaky (>=3.7)", "packaging (>=23)", "pytest (>=7.2.1)", "pytest-env (>=0.8.1)", "pytest-freezegun (>=0.4.2)", "pytest-mock (>=3.10)", "pytest-randomly (>=3.12)", "pytest-timeout (>=2.1)"] 1110 | 1111 | [[package]] 1112 | name = "zipp" 1113 | version = "3.15.0" 1114 | description = "Backport of pathlib-compatible object wrapper for zip files" 1115 | category = "dev" 1116 | optional = false 1117 | python-versions = ">=3.7" 1118 | files = [ 1119 | {file = "zipp-3.15.0-py3-none-any.whl", hash = "sha256:48904fc76a60e542af151aded95726c1a5c34ed43ab4134b597665c86d7ad556"}, 1120 | {file = "zipp-3.15.0.tar.gz", hash = "sha256:112929ad649da941c23de50f356a2b5570c954b65150642bccdd66bf194d224b"}, 1121 | ] 1122 | 1123 | [package.extras] 1124 | docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] 1125 | testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)"] 1126 | 1127 | [metadata] 1128 | lock-version = "2.0" 1129 | python-versions = ">=3.8,<3.12" 1130 | content-hash = "d58bbea0b379f12a18195c9ac4b4db16e916ed188687f74198393ccc22ea20c5" 1131 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "simple-pytree" 3 | version = "0.2.3" 4 | description = "" 5 | authors = ["Cristian Garcia "] 6 | license = "MIT" 7 | readme = "README.md" 8 | packages = [{ include = "simple_pytree" }] 9 | 10 | [tool.poetry.dependencies] 11 | python = ">=3.8" 12 | jax = "*" 13 | jaxlib = "*" 14 | typing-extensions = "*" 15 | 16 | 17 | [tool.poetry.group.dev.dependencies] 18 | pytest = ">=7.2.1" 19 | pytest-cov = ">=4.0.0" 20 | pre-commit = ">=3.0.4" 21 | black = "23.1.0" 22 | isort = "5.12.0" 23 | flax = "*" 24 | 25 | [build-system] 26 | requires = ["poetry-core"] 27 | build-backend = "poetry.core.masonry.api" 28 | 29 | [tool.coverage.report] 30 | exclude_lines = ["@tp.overload"] 31 | -------------------------------------------------------------------------------- /scripts/update_version.py: -------------------------------------------------------------------------------- 1 | import re 2 | from pathlib import Path 3 | 4 | import typer 5 | 6 | 7 | # NOTE: this script could be written bash using sed, but I'm not sure if it's worth it 8 | def main(release_name: str): 9 | release_name = release_name.replace("-create-release", "") 10 | 11 | # Update pyproject.toml 12 | pyproject_path = Path("pyproject.toml") 13 | pyproject_text = pyproject_path.read_text() 14 | pyproject_text = re.sub( 15 | r'version = ".*"', 16 | f'version = "{release_name}"', 17 | pyproject_text, 18 | count=1, 19 | ) 20 | pyproject_path.write_text(pyproject_text) 21 | 22 | # Update __init__.py 23 | init_path = Path("simple_pytree/__init__.py") 24 | init_text = init_path.read_text() 25 | init_text = re.sub( 26 | r'__version__ = "(.*?)"', 27 | f'__version__ = "{release_name}"', 28 | init_text, 29 | count=1, 30 | ) 31 | init_path.write_text(init_text) 32 | 33 | 34 | if __name__ == "__main__": 35 | typer.run(main) 36 | -------------------------------------------------------------------------------- /simple_pytree/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.2.3" 2 | 3 | from .dataclass import dataclass, field, static_field 4 | from .pytree import Pytree, PytreeMeta 5 | 6 | __all__ = ["Pytree", "PytreeMeta", "dataclass", "field", "static_field"] 7 | -------------------------------------------------------------------------------- /simple_pytree/dataclass.py: -------------------------------------------------------------------------------- 1 | import dataclasses 2 | import typing as tp 3 | 4 | import typing_extensions as tpe 5 | 6 | A = tp.TypeVar("A") 7 | 8 | 9 | def field( 10 | *, 11 | default: tp.Any = dataclasses.MISSING, 12 | pytree_node: bool = True, 13 | default_factory: tp.Any = dataclasses.MISSING, 14 | init: bool = True, 15 | repr: bool = True, 16 | hash: tp.Optional[bool] = None, 17 | compare: bool = True, 18 | metadata: tp.Optional[tp.Mapping[str, tp.Any]] = None, 19 | ): 20 | if metadata is None: 21 | metadata = {} 22 | else: 23 | metadata = dict(metadata) 24 | 25 | if "pytree_node" in metadata: 26 | raise ValueError("'pytree_node' found in metadata") 27 | 28 | metadata["pytree_node"] = pytree_node 29 | 30 | return dataclasses.field( # type: ignore 31 | default=default, 32 | default_factory=default_factory, 33 | init=init, 34 | repr=repr, 35 | hash=hash, 36 | compare=compare, 37 | metadata=metadata, 38 | ) 39 | 40 | 41 | def static_field( 42 | *, 43 | default: tp.Any = dataclasses.MISSING, 44 | default_factory: tp.Any = dataclasses.MISSING, 45 | init: bool = True, 46 | repr: bool = True, 47 | hash: tp.Optional[bool] = None, 48 | compare: bool = True, 49 | metadata: tp.Optional[tp.Mapping[str, tp.Any]] = None, 50 | ): 51 | return field( 52 | default=default, 53 | pytree_node=False, 54 | default_factory=default_factory, 55 | init=init, 56 | repr=repr, 57 | hash=hash, 58 | compare=compare, 59 | metadata=metadata, 60 | ) 61 | 62 | 63 | @tp.overload 64 | def dataclass(cls: tp.Type[A]) -> tp.Type[A]: 65 | ... 66 | 67 | 68 | @tp.overload 69 | def dataclass( 70 | *, 71 | init: bool = True, 72 | repr: bool = True, 73 | eq: bool = True, 74 | order: bool = False, 75 | unsafe_hash: bool = False, 76 | frozen: bool = False, 77 | ) -> tp.Callable[[tp.Type[A]], tp.Type[A]]: 78 | ... 79 | 80 | 81 | @tpe.dataclass_transform(field_specifiers=(field, static_field, dataclasses.field)) 82 | def dataclass( 83 | cls: tp.Optional[tp.Type[A]] = None, 84 | init: bool = True, 85 | repr: bool = True, 86 | eq: bool = True, 87 | order: bool = False, 88 | unsafe_hash: bool = False, 89 | frozen: bool = False, 90 | ) -> tp.Union[tp.Type[A], tp.Callable[[tp.Type[A]], tp.Type[A]]]: 91 | decorator = dataclasses.dataclass( 92 | init=init, 93 | repr=repr, 94 | eq=eq, 95 | order=order, 96 | unsafe_hash=unsafe_hash, 97 | frozen=frozen, 98 | ) 99 | 100 | if cls is None: 101 | return decorator 102 | 103 | return decorator(cls) 104 | -------------------------------------------------------------------------------- /simple_pytree/pytree.py: -------------------------------------------------------------------------------- 1 | import dataclasses 2 | import importlib.util 3 | import inspect 4 | import typing as tp 5 | from abc import ABCMeta 6 | from copy import copy 7 | from functools import partial 8 | from types import MappingProxyType 9 | 10 | import jax 11 | 12 | P = tp.TypeVar("P", bound="Pytree") 13 | 14 | 15 | class PytreeMeta(ABCMeta): 16 | def __call__(cls: tp.Type[P], *args: tp.Any, **kwargs: tp.Any) -> P: 17 | obj: P = cls.__new__(cls, *args, **kwargs) 18 | obj.__dict__["_pytree__initializing"] = True 19 | try: 20 | obj.__init__(*args, **kwargs) 21 | finally: 22 | del obj.__dict__["_pytree__initializing"] 23 | 24 | vars_dict = vars(obj) 25 | vars_dict["_pytree__node_fields"] = tuple( 26 | sorted( 27 | field for field in vars_dict if field not in cls._pytree__static_fields 28 | ) 29 | ) 30 | return obj 31 | 32 | 33 | class Pytree(metaclass=PytreeMeta): 34 | _pytree__initializing: bool 35 | _pytree__class_is_mutable: bool 36 | _pytree__static_fields: tp.Tuple[str, ...] 37 | _pytree__node_fields: tp.Tuple[str, ...] 38 | _pytree__setter_descriptors: tp.FrozenSet[str] 39 | 40 | def __init_subclass__(cls, mutable: bool = False): 41 | super().__init_subclass__() 42 | 43 | # gather class info 44 | class_vars = vars(cls) 45 | setter_descriptors = set() 46 | static_fields = _inherited_static_fields(cls) 47 | 48 | # add special static fields 49 | static_fields.add("_pytree__node_fields") 50 | 51 | for field, value in class_vars.items(): 52 | if isinstance(value, dataclasses.Field) and not value.metadata.get( 53 | "pytree_node", True 54 | ): 55 | static_fields.add(field) 56 | 57 | # add setter descriptors 58 | if hasattr(value, "__set__"): 59 | setter_descriptors.add(field) 60 | 61 | static_fields = tuple(sorted(static_fields)) 62 | 63 | # init class variables 64 | cls._pytree__initializing = False 65 | cls._pytree__class_is_mutable = mutable 66 | cls._pytree__static_fields = static_fields 67 | cls._pytree__setter_descriptors = frozenset(setter_descriptors) 68 | 69 | # TODO: clean up this in the future once minimal supported version is 0.4.7 70 | if hasattr(jax.tree_util, "register_pytree_with_keys"): 71 | if ( 72 | "flatten_func" 73 | in inspect.signature(jax.tree_util.register_pytree_with_keys).parameters 74 | ): 75 | jax.tree_util.register_pytree_with_keys( 76 | cls, 77 | partial( 78 | cls._pytree__flatten, 79 | with_key_paths=True, 80 | ), 81 | cls._pytree__unflatten, 82 | flatten_func=partial( 83 | cls._pytree__flatten, 84 | with_key_paths=False, 85 | ), 86 | ) 87 | else: 88 | jax.tree_util.register_pytree_with_keys( 89 | cls, 90 | partial( 91 | cls._pytree__flatten, 92 | with_key_paths=True, 93 | ), 94 | cls._pytree__unflatten, 95 | ) 96 | else: 97 | jax.tree_util.register_pytree_node( 98 | cls, 99 | partial( 100 | cls._pytree__flatten, 101 | with_key_paths=False, 102 | ), 103 | cls._pytree__unflatten, 104 | ) 105 | 106 | # flax serialization support 107 | if importlib.util.find_spec("flax") is not None: 108 | from flax import serialization 109 | 110 | serialization.register_serialization_state( 111 | cls, 112 | partial(cls._to_flax_state_dict, cls._pytree__static_fields), 113 | partial(cls._from_flax_state_dict, cls._pytree__static_fields), 114 | ) 115 | 116 | @classmethod 117 | def _pytree__flatten( 118 | cls, 119 | pytree: "Pytree", 120 | *, 121 | with_key_paths: bool, 122 | ) -> tp.Tuple[tp.Tuple[tp.Any, ...], tp.Mapping[str, tp.Any],]: 123 | all_vars = vars(pytree).copy() 124 | static = {k: all_vars.pop(k) for k in pytree._pytree__static_fields} 125 | 126 | if with_key_paths: 127 | node_values = tuple( 128 | (jax.tree_util.GetAttrKey(field), all_vars.pop(field)) 129 | for field in pytree._pytree__node_fields 130 | ) 131 | else: 132 | node_values = tuple( 133 | all_vars.pop(field) for field in pytree._pytree__node_fields 134 | ) 135 | 136 | if all_vars: 137 | raise ValueError( 138 | f"Unexpected fields in {cls.__name__}: {', '.join(all_vars.keys())}." 139 | "You cannot add new fields to a Pytree after it has been initialized." 140 | ) 141 | 142 | return node_values, MappingProxyType(static) 143 | 144 | @classmethod 145 | def _pytree__unflatten( 146 | cls: tp.Type[P], 147 | static_fields: tp.Mapping[str, tp.Any], 148 | node_values: tp.Tuple[tp.Any, ...], 149 | ) -> P: 150 | pytree = object.__new__(cls) 151 | pytree.__dict__.update(zip(static_fields["_pytree__node_fields"], node_values)) 152 | pytree.__dict__.update(static_fields) 153 | return pytree 154 | 155 | @classmethod 156 | def _to_flax_state_dict( 157 | cls, static_field_names: tp.Tuple[str, ...], pytree: "Pytree" 158 | ) -> tp.Dict[str, tp.Any]: 159 | from flax import serialization 160 | 161 | state_dict = { 162 | name: serialization.to_state_dict(getattr(pytree, name)) 163 | for name in pytree.__dict__ 164 | if name not in static_field_names 165 | } 166 | return state_dict 167 | 168 | @classmethod 169 | def _from_flax_state_dict( 170 | cls, 171 | static_field_names: tp.Tuple[str, ...], 172 | pytree: P, 173 | state: tp.Dict[str, tp.Any], 174 | ) -> P: 175 | """Restore the state of a data class.""" 176 | from flax import serialization 177 | 178 | state = state.copy() # copy the state so we can pop the restored fields. 179 | updates = {} 180 | for name in pytree.__dict__: 181 | if name in static_field_names: 182 | continue 183 | if name not in state: 184 | raise ValueError( 185 | f"Missing field {name} in state dict while restoring" 186 | f" an instance of {type(pytree).__name__}," 187 | f" at path {serialization.current_path()}" 188 | ) 189 | value = getattr(pytree, name) 190 | value_state = state.pop(name) 191 | updates[name] = serialization.from_state_dict(value, value_state, name=name) 192 | if state: 193 | names = ",".join(state.keys()) 194 | raise ValueError( 195 | f'Unknown field(s) "{names}" in state dict while' 196 | f" restoring an instance of {type(pytree).__name__}" 197 | f" at path {serialization.current_path()}" 198 | ) 199 | return pytree.replace(**updates) 200 | 201 | def replace(self: P, **kwargs: tp.Any) -> P: 202 | """ 203 | Replace the values of the fields of the object with the values of the 204 | keyword arguments. If the object is a dataclass, `dataclasses.replace` 205 | will be used. Otherwise, a new object will be created with the same 206 | type as the original object. 207 | """ 208 | if dataclasses.is_dataclass(self): 209 | return dataclasses.replace(self, **kwargs) 210 | 211 | unknown_keys = set(kwargs) - set(vars(self)) 212 | if unknown_keys: 213 | raise ValueError( 214 | f"Trying to replace unknown fields {unknown_keys} " 215 | f"for '{type(self).__name__}'" 216 | ) 217 | 218 | pytree = copy(self) 219 | pytree.__dict__.update(kwargs) 220 | return pytree 221 | 222 | if not tp.TYPE_CHECKING: 223 | 224 | def __setattr__(self: P, field: str, value: tp.Any): 225 | if self._pytree__initializing or field in self._pytree__setter_descriptors: 226 | pass 227 | elif not hasattr(self, field) and not self._pytree__initializing: 228 | raise AttributeError( 229 | f"Cannot add new fields to {type(self)} after initialization" 230 | ) 231 | elif not self._pytree__class_is_mutable: 232 | raise AttributeError( 233 | f"{type(self)} is immutable, trying to update field {field}" 234 | ) 235 | 236 | object.__setattr__(self, field, value) 237 | 238 | 239 | def _inherited_static_fields(cls: type) -> tp.Set[str]: 240 | static_fields = set() 241 | for parent_class in cls.mro(): 242 | if parent_class is not cls and parent_class is not Pytree: 243 | if issubclass(parent_class, Pytree): 244 | static_fields.update(parent_class._pytree__static_fields) 245 | elif dataclasses.is_dataclass(parent_class): 246 | for field in dataclasses.fields(parent_class): 247 | if not field.metadata.get("pytree_node", True): 248 | static_fields.add(field.name) 249 | return static_fields 250 | -------------------------------------------------------------------------------- /tests/test_pytree.py: -------------------------------------------------------------------------------- 1 | import dataclasses 2 | from typing import Generic, TypeVar 3 | 4 | import jax 5 | import pytest 6 | from flax import serialization 7 | 8 | from simple_pytree import Pytree, dataclass, field, static_field 9 | 10 | 11 | class TestPytree: 12 | def test_immutable_pytree(self): 13 | class Foo(Pytree): 14 | x: int = static_field() 15 | 16 | def __init__(self, y) -> None: 17 | self.x = 2 18 | self.y = y 19 | 20 | pytree = Foo(y=3) 21 | 22 | leaves = jax.tree_util.tree_leaves(pytree) 23 | assert leaves == [3] 24 | 25 | pytree = jax.tree_map(lambda x: x * 2, pytree) 26 | assert pytree.x == 2 27 | assert pytree.y == 6 28 | 29 | pytree = pytree.replace(x=3) 30 | assert pytree.x == 3 31 | assert pytree.y == 6 32 | 33 | with pytest.raises( 34 | AttributeError, match="is immutable, trying to update field" 35 | ): 36 | pytree.x = 4 37 | 38 | def test_immutable_pytree_dataclass(self): 39 | @dataclass(frozen=True) 40 | class Foo(Pytree): 41 | y: int = field() 42 | x: int = static_field(default=2) 43 | 44 | pytree = Foo(y=3) 45 | 46 | leaves = jax.tree_util.tree_leaves(pytree) 47 | assert leaves == [3] 48 | 49 | pytree = jax.tree_map(lambda x: x * 2, pytree) 50 | assert pytree.x == 2 51 | assert pytree.y == 6 52 | 53 | pytree = pytree.replace(x=3) 54 | assert pytree.x == 3 55 | assert pytree.y == 6 56 | 57 | with pytest.raises(AttributeError, match="cannot assign to field"): 58 | pytree.x = 4 59 | 60 | def test_jit(self): 61 | @dataclass 62 | class Foo(Pytree): 63 | a: int 64 | b: int = static_field() 65 | 66 | module = Foo(a=1, b=2) 67 | 68 | @jax.jit 69 | def f(m: Foo): 70 | return m.a + m.b 71 | 72 | assert f(module) == 3 73 | 74 | def test_flax_serialization(self): 75 | class Bar(Pytree): 76 | a: int = static_field() 77 | 78 | def __init__(self, a, b): 79 | self.a = a 80 | self.b = b 81 | 82 | @dataclass 83 | class Foo(Pytree): 84 | bar: Bar 85 | c: int 86 | d: int = static_field() 87 | 88 | foo: Foo = Foo(bar=Bar(a=1, b=2), c=3, d=4) 89 | 90 | state_dict = serialization.to_state_dict(foo) 91 | 92 | assert state_dict == { 93 | "bar": { 94 | "b": 2, 95 | }, 96 | "c": 3, 97 | } 98 | 99 | state_dict["bar"]["b"] = 5 100 | 101 | foo = serialization.from_state_dict(foo, state_dict) 102 | 103 | assert foo.bar.b == 5 104 | 105 | del state_dict["bar"]["b"] 106 | 107 | with pytest.raises(ValueError, match="Missing field"): 108 | serialization.from_state_dict(foo, state_dict) 109 | 110 | state_dict["bar"]["b"] = 5 111 | 112 | # add unknown field 113 | state_dict["x"] = 6 114 | 115 | with pytest.raises(ValueError, match="Unknown field"): 116 | serialization.from_state_dict(foo, state_dict) 117 | 118 | def test_generics(self): 119 | T = TypeVar("T") 120 | 121 | class MyClass(Pytree, Generic[T]): 122 | def __init__(self, x: T): 123 | self.x = x 124 | 125 | MyClass[int] 126 | 127 | def test_key_paths(self): 128 | @dataclass 129 | class Bar(Pytree): 130 | a: int = 1 131 | b: int = static_field(default=2) 132 | 133 | @dataclass 134 | class Foo(Pytree): 135 | x: int = 3 136 | y: int = static_field(default=4) 137 | z: Bar = field(default_factory=Bar) 138 | 139 | foo = Foo() 140 | 141 | path_values, treedef = jax.tree_util.tree_flatten_with_path(foo) 142 | path_values = [(list(map(str, path)), value) for path, value in path_values] 143 | 144 | assert path_values[0] == ([".x"], 3) 145 | assert path_values[1] == ([".z", ".a"], 1) 146 | 147 | def test_setter_attribute_allowed(self): 148 | n = None 149 | 150 | class SetterDescriptor: 151 | def __set__(self, _, value): 152 | nonlocal n 153 | n = value 154 | 155 | class Foo(Pytree): 156 | x = SetterDescriptor() 157 | 158 | foo = Foo() 159 | foo.x = 1 160 | 161 | assert n == 1 162 | 163 | def test_replace_unknown_fields_error(self): 164 | class Foo(Pytree): 165 | pass 166 | 167 | with pytest.raises(ValueError, match="Trying to replace unknown fields"): 168 | Foo().replace(y=1) 169 | 170 | def test_dataclass_inheritance(self): 171 | @dataclass 172 | class A(Pytree): 173 | a: int = 1 174 | b: int = static_field(default=2) 175 | 176 | @dataclass 177 | class B(A): 178 | c: int = 3 179 | 180 | pytree = B() 181 | leaves = jax.tree_util.tree_leaves(pytree) 182 | assert leaves == [1, 3] 183 | 184 | def test_pytree_with_new(self): 185 | class A(Pytree): 186 | def __init__(self, a): 187 | self.a = a 188 | 189 | def __new__(cls, a): 190 | return super().__new__(cls) 191 | 192 | pytree = A(a=1) 193 | 194 | pytree = jax.tree_map(lambda x: x * 2, pytree) 195 | 196 | def test_deterministic_order(self): 197 | class A(Pytree): 198 | def __init__(self, order: bool): 199 | if order: 200 | self.a = 1 201 | self.b = 2 202 | else: 203 | self.b = 2 204 | self.a = 1 205 | 206 | p1 = A(order=True) 207 | p2 = A(order=False) 208 | 209 | leaves1 = jax.tree_util.tree_leaves(p1) 210 | leaves2 = jax.tree_util.tree_leaves(p2) 211 | 212 | assert leaves1 == leaves2 213 | 214 | 215 | class TestMutablePytree: 216 | def test_pytree(self): 217 | class Foo(Pytree, mutable=True): 218 | x: int = static_field() 219 | 220 | def __init__(self, y) -> None: 221 | self.x = 2 222 | self.y = y 223 | 224 | pytree = Foo(y=3) 225 | 226 | leaves = jax.tree_util.tree_leaves(pytree) 227 | assert leaves == [3] 228 | 229 | pytree = jax.tree_map(lambda x: x * 2, pytree) 230 | assert pytree.x == 2 231 | assert pytree.y == 6 232 | 233 | pytree = pytree.replace(x=3) 234 | assert pytree.x == 3 235 | assert pytree.y == 6 236 | 237 | # test mutation 238 | pytree.x = 4 239 | assert pytree.x == 4 240 | 241 | def test_no_new_fields_after_init(self): 242 | class Foo(Pytree, mutable=True): 243 | def __init__(self, x): 244 | self.x = x 245 | 246 | foo = Foo(x=1) 247 | foo.x = 2 248 | 249 | with pytest.raises(AttributeError, match=r"Cannot add new fields to"): 250 | foo.y = 2 251 | 252 | def test_pytree_dataclass(self): 253 | @dataclass 254 | class Foo(Pytree, mutable=True): 255 | y: int = field() 256 | x: int = static_field(default=2) 257 | 258 | pytree: Foo = Foo(y=3) 259 | 260 | leaves = jax.tree_util.tree_leaves(pytree) 261 | assert leaves == [3] 262 | 263 | pytree = jax.tree_map(lambda x: x * 2, pytree) 264 | assert pytree.x == 2 265 | assert pytree.y == 6 266 | 267 | pytree = pytree.replace(x=3) 268 | assert pytree.x == 3 269 | assert pytree.y == 6 270 | 271 | # test mutation 272 | pytree.x = 4 273 | assert pytree.x == 4 274 | 275 | def test_dataclass_inheritance(self): 276 | A = dataclasses.make_dataclass( 277 | "A", 278 | [("x", int), "y", ("z", int, static_field(default=5))], 279 | ) 280 | 281 | @dataclass 282 | class B(Pytree, A): 283 | ... 284 | 285 | b = B(1, 2) 286 | 287 | assert b.x == 1 288 | assert b.y == 2 289 | assert b.z == 5 290 | 291 | assert jax.tree_util.tree_leaves(b) == [1, 2] 292 | --------------------------------------------------------------------------------