From b5c25822ff2428c8b0db02dca1c9aab5340e2d08 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Mon, 11 Jul 2022 13:03:39 +0100 Subject: [PATCH] Fixes (#1) * Fixes * Black * Test fix * Test fix * workflow fixes --- .github/workflows/release.yml | 2 +- .github/workflows/run_tests.yml | 2 +- .isort.cfg | 1 + jaxtyping/__init__.py | 2 +- jaxtyping/array_types.py | 3 ++- jaxtyping/import_hook.py | 2 +- jaxtyping/pytree_type.py | 22 +++++++++++++++++++++- setup.py | 2 +- 8 files changed, 29 insertions(+), 7 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 75e31fe..ad9c799 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -33,7 +33,7 @@ jobs: with: python-version: "3.8" test-script: | - python -m pip install pytest jax jaxlib typeguard + python -m pip install pytest beartype equinox jaxlib cp -r ${{ github.workspace }}/test ./test pytest pypi-token: ${{ secrets.pypi_token }} diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 31a75f2..56d0d23 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -42,7 +42,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install pytest wheel jaxlib + python -m pip install pytest wheel beartype equinox jaxlib - name: Checks with pre-commit uses: pre-commit/action@v2.0.3 diff --git a/.isort.cfg b/.isort.cfg index 64737cc..bdc8fe4 100644 --- a/.isort.cfg +++ b/.isort.cfg @@ -3,3 +3,4 @@ force_alphabetical_sort_within_sections=true lines_after_imports=2 profile=black treat_comments_as_code=true +extra_standard_library=typing_extensions diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 8547022..bf4a5ad 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -51,4 +51,4 @@ from .import_hook import install_import_hook from .pytree_type import PyTree -__version__ = "0.0.1" +__version__ = "0.0.2" diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 667a090..348ee06 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -18,7 +18,8 @@ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. import functools as ft -from typing import Any, Dict, List, Literal, NoReturn, Optional, Tuple, Union +from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union +from typing_extensions import Literal import jax.numpy as jnp diff --git a/jaxtyping/import_hook.py b/jaxtyping/import_hook.py index a419fde..31f74e8 100644 --- a/jaxtyping/import_hook.py +++ b/jaxtyping/import_hook.py @@ -104,7 +104,7 @@ class _JaxtypingTransformer(ast.NodeVisitor): 0, ast.Attribute( ast.Name(id="jaxtyping", ctx=ast.Load()), "jaxtyped", ast.Load() - ) + ), ) if self._typechecker is not None: # Place at the end of the decorator list, as decorators diff --git a/jaxtyping/pytree_type.py b/jaxtyping/pytree_type.py index 7a7d894..c19a0bc 100644 --- a/jaxtyping/pytree_type.py +++ b/jaxtyping/pytree_type.py @@ -18,11 +18,27 @@ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. import functools as ft +from typing import Generic, TypeVar import jax import typeguard +_T = TypeVar("_T") + + +class _FakePyTree(Generic[_T]): + pass + + +_FakePyTree.__name__ = "PyTree" +_FakePyTree.__qualname__ = "PyTree" +# Can't do type("PyTree", (Generic[_T],), {}) because dynamic subclassing of typeforms +# isn't allowed. +# Can't do types.new_class("PyTree", (Generic[_T],), {}) because that has __module__ +# "types", e.g. we get types.PyTree[int]. + + class _MetaPyTree(type): def __call__(self, *args, **kwargs): raise RuntimeError("PyTree cannot be instantiated") @@ -32,7 +48,8 @@ class _MetaPyTree(type): @ft.lru_cache(maxsize=None) def __getitem__(cls, item): - return _MetaSubscriptPyTree(f"PyTree[{item.__name__}]", (), {"leaftype": item}) + name = str(_FakePyTree[item]) + return _MetaSubscriptPyTree(name, (), {"leaftype": item}) class _MetaSubscriptPyTree(type): @@ -63,3 +80,6 @@ class _MetaSubscriptPyTree(type): PyTree = _MetaPyTree("PyTree", (), {}) +# Can't do `class PyTree(Generic[_T]): ...` because we need to override the +# instancecheck for PyTree[foo], but we subclassing +# `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed. diff --git a/setup.py b/setup.py index 9776fcd..1a97619 100644 --- a/setup.py +++ b/setup.py @@ -63,7 +63,7 @@ python_requires = "~=3.7" # We use typeguard internally (in a fairly minimal way), but it's not required that # end users make the same choice. -install_requires = ["jax>=0.3.4", "typeguard>=2.13.3"] +install_requires = ["jax>=0.3.4", "typeguard>=2.13.3", "typing_extensions>=4.2.0"] entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"])