├── .coveragerc ├── .flake8 ├── .github └── workflows │ ├── build.yml │ ├── coverage.yml │ ├── lint.yml │ ├── mypy.yml │ └── publish.yml ├── .gitignore ├── LICENSE ├── README.md ├── deprecated.md ├── jax_dataclasses ├── __init__.py ├── _copy_and_mutate.py ├── _dataclasses.py ├── _enforced_annotations.py ├── _get_type_hints.py ├── _jit.py └── py.typed ├── mypy.ini ├── setup.py └── tests ├── conftest.py ├── test_annotated_arrays.py ├── test_copy_and_mutate.py ├── test_cycle.py ├── test_dataclass.py ├── test_jit_ignore_py37.py ├── test_serialization.py ├── test_variadic_generic_py312.py └── test_vmap.py /.coveragerc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/.coveragerc -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/.flake8 -------------------------------------------------------------------------------- /.github/workflows/build.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/.github/workflows/build.yml -------------------------------------------------------------------------------- /.github/workflows/coverage.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/.github/workflows/coverage.yml -------------------------------------------------------------------------------- /.github/workflows/lint.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/.github/workflows/lint.yml -------------------------------------------------------------------------------- /.github/workflows/mypy.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/.github/workflows/mypy.yml -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/.github/workflows/publish.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/README.md -------------------------------------------------------------------------------- /deprecated.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/deprecated.md -------------------------------------------------------------------------------- /jax_dataclasses/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/jax_dataclasses/__init__.py -------------------------------------------------------------------------------- /jax_dataclasses/_copy_and_mutate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/jax_dataclasses/_copy_and_mutate.py -------------------------------------------------------------------------------- /jax_dataclasses/_dataclasses.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/jax_dataclasses/_dataclasses.py -------------------------------------------------------------------------------- /jax_dataclasses/_enforced_annotations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/jax_dataclasses/_enforced_annotations.py -------------------------------------------------------------------------------- /jax_dataclasses/_get_type_hints.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/jax_dataclasses/_get_type_hints.py -------------------------------------------------------------------------------- /jax_dataclasses/_jit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/jax_dataclasses/_jit.py -------------------------------------------------------------------------------- /jax_dataclasses/py.typed: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- 1 | [mypy] 2 | python_version = 3.12 3 | ignore_missing_imports = True 4 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/setup.py -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/conftest.py -------------------------------------------------------------------------------- /tests/test_annotated_arrays.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/test_annotated_arrays.py -------------------------------------------------------------------------------- /tests/test_copy_and_mutate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/test_copy_and_mutate.py -------------------------------------------------------------------------------- /tests/test_cycle.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/test_cycle.py -------------------------------------------------------------------------------- /tests/test_dataclass.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/test_dataclass.py -------------------------------------------------------------------------------- /tests/test_jit_ignore_py37.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/test_jit_ignore_py37.py -------------------------------------------------------------------------------- /tests/test_serialization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/test_serialization.py -------------------------------------------------------------------------------- /tests/test_variadic_generic_py312.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/test_variadic_generic_py312.py -------------------------------------------------------------------------------- /tests/test_vmap.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brentyi/jax_dataclasses/HEAD/tests/test_vmap.py --------------------------------------------------------------------------------