commit 7ac6ee04a8ec2f1a6b724a1ed2414d438069f2cf Author: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Fri Jun 24 17:02:18 2022 +0000 initial commit diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..7b18abd --- /dev/null +++ b/.flake8 @@ -0,0 +1,4 @@ +[flake8] +max-line-length = 120 +ignore = W291,W293,W503,W504,E123,E126,E203,E402,E701,E731,F722 +per-file-ignores = __init__.py: F401 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..75e31fe --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,45 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +name: Release + +on: + push: + branches: + - main + +jobs: + build: + runs-on: ubuntu-latest + steps: + - name: Release + uses: patrick-kidger/action_update_python_project@v1 + with: + python-version: "3.8" + test-script: | + python -m pip install pytest jax jaxlib typeguard + cp -r ${{ github.workspace }}/test ./test + pytest + pypi-token: ${{ secrets.pypi_token }} + github-user: patrick-kidger + github-token: ${{ github.token }} + email-user: ${{ secrets.email_user }} + email-token: ${{ secrets.email_token }} + email-server: ${{ secrets.email_server }} + email-target: ${{ secrets.email_target }} diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml new file mode 100644 index 0000000..31a75f2 --- /dev/null +++ b/.github/workflows/run_tests.yml @@ -0,0 +1,53 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +name: Run tests + +on: + pull_request: + +jobs: + run-tests: + strategy: + matrix: + python-version: [ 3.7, 3.8, 3.9 ] + os: [ ubuntu-latest ] + fail-fast: false + runs-on: ${{ matrix.os }} + steps: + - name: Checkout code + uses: actions/checkout@v2 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install pytest wheel jaxlib + + - name: Checks with pre-commit + uses: pre-commit/action@v2.0.3 + + - name: Test with pytest + run: | + python -m pip install . + python -m pytest --durations=0 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b1b6a31 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +**/__pycache__ +*.egg-info + diff --git a/.isort.cfg b/.isort.cfg new file mode 100644 index 0000000..64737cc --- /dev/null +++ b/.isort.cfg @@ -0,0 +1,5 @@ +[settings] +force_alphabetical_sort_within_sections=true +lines_after_imports=2 +profile=black +treat_comments_as_code=true diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..e975f15 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,38 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +repos: + - repo: https://github.com/ambv/black + rev: 22.3.0 + hooks: + - id: black + - repo: https://github.com/nbQA-dev/nbQA + rev: 1.2.3 + hooks: + - id: nbqa-black + - id: nbqa-isort + - id: nbqa-flake8 + - repo: https://github.com/PyCQA/isort + rev: 5.10.1 + hooks: + - id: isort + - repo: https://github.com/pycqa/flake8 + rev: 4.0.1 + hooks: + - id: flake8 diff --git a/API.md b/API.md new file mode 100644 index 0000000..bcf8ce1 --- /dev/null +++ b/API.md @@ -0,0 +1,196 @@ +# Full API + +## Annotating array types + +Each array is denoted by a type `dtype[shape]`, such as `f32["batch channels"]`. + +### Shape + +The shape should be a string of space-separated symbols, such as "a b c d". Each symbol can be: +- `int`: fixed-size axis, e.g. `f32["28 28"]`. +- `str`: variable-size axis, e.g. `f32["channels"]`. +- `_`: anonymous axis, e.g. `f32["batch channels _ _"]`. +- `...`: anonymous zero or more axes, e.g. `f32["... c h w"]` +- `*name`: zero or more variable-size axes, e.g. `f32["*batch c h w"]` +- Append `#` to a dimension size to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable. + +When calling a function, variable-size axes will be matched up across all arguments and checked for consistency. (See [runtime type checking](#runtime-type-checking) below.) + +Some notes: +- To denote a scalar shape use `""`, e.g. `f32[""]`. +- To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `f32["..."]`. +- You cannot have multiple variadic axes, i.e. you can only use `...` or `*name` at most once in each array. +- An example of broadcasting in one dimension: `add(x: f32["foo#"], y: f32["foo#"]) -> f32["foo#"]`. +- An example of broadcasting multiple dimensions: `add(x: f32["*foo#"], y: f32["*foo#"]) -> f32["*foo#"]`. + +### Dtype + +The dtype should be any one of (imported from `jaxtyping`): +- Any dtype at all: `Array` + - Boolean: `b` + - Any integer, unsigned integer, floating, or complex: `n` (for number) + - Any floating or complex: `x` (for inexact) + - Any floating point: `f` + - Floating point: `bf16`, `f16`, `f32`, `f64` (`bf16` is bfloat16) + - Any complex: `c` + - Complexes: `c64`, `c128` + - Any integer or unsigned intger: `t` (for integer) + - Any unsigned integer: `u` + - Unsigned integer: `u8`, `u16`, `u32`, `u64` + - Any signed integer: `i` + - Signed integer: `i8`, `i16`, `i32`, `i64` + +Unless you really want to force a particular precision, then for most applications you should probably allow any floating-point, any integer, etc. That is, use +```python +from jaxtyping import f +f["some_shape"] +``` +rather than +```python +from jaxtyping import f32 +f32["some_shape"] +``` + +## PyTrees + +### `jaxtyping.PyTree` + +Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or `PyTree[Union[str, f32["b c"]]]`. + +You can leave off the `[...]`, in which case `PyTree` is simply a suggestively-named alternative to `Any`. ([By definition all types are PyTrees.](https://jax.readthedocs.io/en/latest/pytrees.html)) + +## Runtime type checking + +Single-argument type checking will work with any runtime type checker out-of-the-box. + +To enable multi-argument consistency checks (i.e. that shapes match up between arrays), then you have two options, as discussed below. (And if either are too much magic for you, you can safely use neither and stick to just single-argument type checking.) + +Regardless of your choice, **this approach synergises beautifully with `jax.jit`!** All shape checks will be performed at trace-time only, and will not impact runtime performance. + +### `jaxtyping.jaxtyped` + +Decorate a function with this to have shapes checked for consistency across multiple arguments. + +Example: + +```python +# Import both the annotation and the `jaxtyped` decorator from `jaxtyping` +from jaxtyping import f32, jaxtyped + +# Use your favourite typechecker: usually one of the two lines below. +from typeguard import typechecked as typechecker +from beartype import beartype as typechecker + +# Write your function. @jaxtyped must be applied above @typechecker! +@jaxtyped +@typechecker +def batch_outer_product(x: f32["b c1"], y: f32["b c2"]) -> f32["b c1 c2"]: + return x[:, :, None] * y[:, None, :] +``` + +Note that `@jaxtyped` is applied above the type checker. + +#### `jaxtyping.jaxtyped` for advanced users + +Put precisely, all `isinstance` shape checks are scoped to the thread-local dynamic context +of a `jaxtyped` call. A new dynamic context will allow different dimensions +sizes to be bound to the same name. After this new dynamic context is finished +then the old one is returned to. + +For example, this means you could leave off the `@jaxtyped` decorator to enforce that +this function use the same axes sizes as the function it was called from. + +Likewise, this means you can use `isinstance` checks inside a function body +and have them contribute to the same collection of consistency checks performed +by a typechecker against its arguments. (Or even forgo a typechecker altogether, +and just do your own manual `isinstance` checks.) + +Only `isinstance` checks that pass will contribute to the store of axis name-size pairs; those +that fail will not. As such it is safe to write e.g. `assert not isinstance(x, +f32["foo"])`. + +### `jaxtyping.install_import_hook` + +It can be a lot of effort to add `@jaxtyped` decorators all over your codebase. +(Not to mention that double-decorators everywhere are a bit ugly.) The easier +option is usually to use the import import hook. + +Example: + +```python +from jaxtyping import install_import_hook +# Plus either one of the following: +install_import_hook("foo", ("typeguard", "typechecked")) # decorate @jaxtyped and @typeguard.typechecked +install_import_hook("foo", ("beartype", "beartype")) # decorate @jaxtyped and @beartype.beartype +install_import_hook("foo", None) # decorate only @jaxtyped (if you have manually applied typechecking decorators) +``` + +Any module imported **afterwards**, whose name begins with the specified string, will automatically have both `@jaxtyped` and the specified typechecker applied to all of their functions. (E.g. in the above example `foo`, `foo.bar`, `foo.bar.qux` would all be hook'd). + +The import hook may be uninstalled after you've imported all the modules you're interested in: +```python +hook = install_import_hook(...) +... # perform imports +hook.uninstall() +``` + +The import hook can be applied to multiple packages via +```python +install_import_hook(["foo", "bar.baz"], ...) +``` + +**Example: writing an end-user script** + +```python +### entry_point.py +from jaxtyping import install_import_hook +install_import_hook("do_stuff", ("typeguard", "typechecked")) +import do_stuff + +### do_stuff.py +from jaxtyping import f32 + +def g(x: f32["..."]): + ... +``` + +**Example: writing a library** + +```python +### __init__.py +from jaxtyping import install_import_hook +hook = install_import_hook("my_library_name", ("beartype", "beartype")) +from .subpackage import foo # full name is my_library_name.subpackage so will be hook'd +from .another_subpackage import bar # full name is my_library_name.another_subpackage so will be hook'd. +hook.uninstall() +del hook, install_import_hook, jaxtyping # keep interface tidy +``` + +#### pytest hook + +The import hook can be installed at test-time only, as a pytest hook. The syntax is +``` +pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype +``` +which will apply the import hook to all modules whose names start with either `foo` or `bar.baz`. The typechecker used in this example is `beartype.beartype`. + +## Abstract base classes + +### `jaxtyping.AbstractDtype` + +The base class of all dtypes. This can be used to create your own custom collection of dtypes (analogous to `n`, `x` etc.) For example: +```python +class u8_or_u16(AbstractDtype): + dtypes = ["uint8", "uint16"] + +u8_or_u16["shape"] +``` +which is functionally equivalent to +```python +Union[u8["shape"], u16["shape"]] +``` + +### `jaxtyping.AbstractArray` + +The base class of all shape-and-dtype-specified arrays, e.g. it's a base class +for `f32["foo"]`. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..79957ad --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,54 @@ +# Contributing + +Contributions (pull requests) are very welcome! Here's how to get started. + +--- + +First fork the library on GitHub. + +Then clone and install the library in development mode: + +```bash +git clone https://github.com/your-username-here/jaxtyping.git +cd jaxtyping +pip install -e . +``` + +Then install the pre-commit hook: + +```bash +pip install pre-commit +pre-commit install +``` + +These hooks use Black and isort to format the code, and flake8 to lint it. + +Now make your changes. Make sure to include additional tests if necessary. + +Next verify the tests all pass: + +```bash +pip install pytest +pytest +``` + +Then push your changes back to your fork of the repository: + +```bash +git push +``` + +Finally, open a pull request on GitHub! + +## Contributor License Agreement + +Contributions to this project must be accompanied by a Contributor License +Agreement (CLA). You (or your employer) retain the copyright to your +contribution; this simply gives us permission to use and redistribute your +contributions as part of the project. Head over to + to see your current agreements on file or +to sign a new one. + +You generally only need to submit a CLA once, so if you've already submitted one +(even if it was for a different project), you probably don't need to do it +again. diff --git a/FAQ.md b/FAQ.md new file mode 100644 index 0000000..1c73509 --- /dev/null +++ b/FAQ.md @@ -0,0 +1,43 @@ +# FAQ + +## `flake8` is throwing an error. + +In type annotations, strings are used for two different things. Sometimes they're strings. Sometimes they're "forward references", used to refer to a type that will be defined later. + +Some tooling in the Python ecosystem assumes that only the latter is true, and will throw spurious errors if you try to use a string just as a string (like we do). + +In the case of `flake8`, at least, this is easily resolved. Multi-dimensional arrays (e.g. `f32["b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `f32["x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `f32[" x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before. + +## What about support for static type checkers, like `mypy`, `pyright`, etc.? + +Nope. + +Python's static typing ecosystem is a complicated collection of edge cases. Many of them block ML/scientific computing in particular. A few examples: + +1. The static type system is intrinsically not expressive enough to describe operations like concatenation, stacking, or broadcasting. + +2. Axes have to be lifted to type-level variables. Meanwhile the approach taken in libraries like `jaxtyping` and [TorchTyping](https://github.com/patrick-kidger/torchtyping) is to use value-level variables for types: because that's what the underlying JAX, PyTorch etc. libraries use! As such, making a static type checker work with these libraries would require either fundamentally rewriting these libraries, or exhaustively maintaining type stubs for them, and would *still* require a `typing.cast` any time you use anything unstubbed (e.g. any third party library, or part of your codebase you haven't typed yet). This is a huge maintenance burden for anyone. + +3. Static type checkers have a variety of bugs that affect this use case. `mypy` doesn't support `Protocol`s correctly. `pyright` doesn't support genericised subprotocols. etc. + +4. Variadic generics exist. Variadic protocols do not. (It's not clear that these have been contemplated.) + +5. The syntax for static typing is verbose. You have to write things like `Array[Unpack[AnyShape], Literal[3], Height, Width]` instead of `Array["... 3 height width"]`. + +6. [The underlying type system has flaws](https://github.com/patrick-kidger/torchtyping/issues/37#issuecomment-1153294196). [The numeric tower is broken](https://stackoverflow.com/a/69383462); [int is not a number](https://github.com/python/mypy/issues/3186#issuecomment-885718629); [virtual base classes don't work](https://github.com/python/mypy/issues/2922); [complex lies about having comparison operations, so type checkers have to lie about that lie in order to remove them again](https://posita.github.io/numerary/0.4/whytho/); `typing.*` don't work with `isinstance`; co/contra-variance are baked into containers (not specified at use-time); `dict` is variadic despite... not being variadic; bool is a subclass of int (!); ... etc. etc. + +## What about [PEP 646](https://www.python.org/dev/peps/pep-0646/) and variadic generics? + +[Doesn't change the previous issues, unfortunately.](https://github.com/patrick-kidger/torchtyping/issues/37) All the problems of the previous heading still hold true. They're just also true for types like `AnyDimensionalArray[Batch, Channels, AsManyArgumentsAsWePlease]` as well as types like `TwoDimensionalArray[Batch, Channels]`. + +## Is the lack of interaction with static typing a problem? + +At least for any software that is mostly just running JAX code, no! + +The correct way to use JAX is to put together all your operations, and then put a single `jax.jit` right at the very top. This gives you optimal speed; anything else will be unnecessarily (and substantially) slower. + +This means that all the type checking only gets resolved once: at trace time. Afterwards JAX still lowers everything down to the same optimised code. + +In some sense, `python myprogram.py` just ends up doing the same as `mypy myprogram.py`. Except instead of throwing away all the work used to parse your code, build the abstract syntax tree, etc. (and requiring you to then run `python myprogram.py` afterwards to actually use it), it can keep it around and just run your code immediately. + +TL;DR: `jax.jit` is amazing. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..43b5177 --- /dev/null +++ b/LICENSE @@ -0,0 +1,51 @@ +MIT License + +Copyright (c) 2022 Google LLC + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + + + + +--- +Sections of the code were modified from https://github.com/agronholm/typeguard +under the terms of the MIT license, reproduced below. +--- + +MIT License + +Copyright (c) Alex Grönholm + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..dd2678e --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,2 @@ +include LICENSE +prune tests diff --git a/README.md b/README.md new file mode 100644 index 0000000..9bc0e93 --- /dev/null +++ b/README.md @@ -0,0 +1,56 @@ +

jaxtyping

+ +Type annotations **and runtime checking** for: + +1. shape and dtype of [JAX](https://github.com/google/jax) arrays; +2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html). + +**For example:** +```python +from jaxtyping import f32, PyTree + +def matrix_multiply(x: f32["dim1 dim2"], y: f32["dim2 dim3"]) -> f32["dim1 dim3"]: + ... + +def accepts_pytree_of_ints(x: PyTree[int]): + ... + +def accepts_pytree_of_arrays(x: PyTree[f32["batch c1 c2"]]): + ... +``` + +## Installation + +```bash +pip install jaxtyping +``` + +Requires JAX 0.3.4+. + +Also install your favourite runtime type-checking package. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments). + +## Documentation + +[Full API reference](./API.md) + +[FAQ (static type checking, flake8, etc.)](./FAQ.md) + +## Finally + +### See also: other tools in the JAX ecosystem + +Neural networks: [Equinox](https://github.com/patrick-kidger/equinox). + +Numerical differential equation solvers: [Diffrax](https://github.com/patrick-kidger/diffrax). + +SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2jax](https://github.com/google/sympy2jax). + +### Acknowledgements + +Shape annotations + runtime type checking is inspired by [TorchTyping](https://github.com/patrick-kidger/torchtyping). + +The concise syntax is inspired by [etils.array_types](https://github.com/google/etils/tree/main/etils/array_types). + +### Disclaimer + +This is not an official Google product. diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py new file mode 100644 index 0000000..8547022 --- /dev/null +++ b/jaxtyping/__init__.py @@ -0,0 +1,54 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +from .array_types import ( + AbstractArray, + AbstractDtype, + Array, + b, + bf16, + c, + c64, + c128, + f, + f16, + f32, + f64, + get_array_name_format, + i, + i8, + i16, + i32, + i64, + n, + set_array_name_format, + t, + u, + u8, + u16, + u32, + u64, + x, +) +from .decorator import jaxtyped +from .import_hook import install_import_hook +from .pytree_type import PyTree + + +__version__ = "0.0.1" diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py new file mode 100644 index 0000000..667a090 --- /dev/null +++ b/jaxtyping/array_types.py @@ -0,0 +1,314 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# 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 + +import jax.numpy as jnp + +from .decorator import storage + + +_array_name_format = "dtype_and_shape" + + +def get_array_name_format(): + return _array_name_format + + +def set_array_name_format(value): + global _array_name_format + _array_name_format = value + + +_any_dtype = object() + +_anonymous_dim = object() +_anonymous_variadic_dim = object() + + +class _NamedDim: + def __init__(self, name, broadcastable): + self.name = name + self.broadcastable = broadcastable + + +class _NamedVariadicDim: + def __init__(self, name, broadcastable): + self.name = name + self.broadcastable = broadcastable + + +class _FixedDim: + def __init__(self, size, broadcastable): + self.size = size + self.broadcastable = broadcastable + + +_AbstractDimOrVariadicDim = Union[ + Literal[_anonymous_dim], + Literal[_anonymous_variadic_dim], + _NamedDim, + _NamedVariadicDim, + _FixedDim, +] +_AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim] + + +def _check_dims( + cls_dims: List[_AbstractDim], + obj_shape: Tuple[int], + memo: Dict[str, Union[int, Tuple[int]]], +): + assert len(cls_dims) == len(obj_shape) + for cls_dim, obj_size in zip(cls_dims, obj_shape): + if cls_dim is _anonymous_dim: + pass + elif cls_dim.broadcastable and obj_size == 1: + pass + elif type(cls_dim) is _FixedDim: + if cls_dim.size != obj_size: + return False + else: + assert type(cls_dim) is _NamedDim + try: + cls_size = memo[cls_dim.name] + except KeyError: + memo[cls_dim.name] = obj_size + else: + if cls_size != obj_size: + return False + return True + + +class _MetaAbstractArray(type): + def __instancecheck__(cls, obj): + if not isinstance(obj, jnp.ndarray): + return False + + if cls.dtypes is not _any_dtype and obj.dtype not in cls.dtypes: + return False + + if len(storage.memo_stack) == 0: + # `isinstance` happening outside any @jaxtyped decorators, e.g. at the + # global scope. In this case just create a temporary memo, since we're not + # going to be comparing against any stored values anyway. + memo = {} + temp_memo = True + else: + # Make a copy so we don't mutate the original memo during the shape check. + memo = storage.memo_stack[-1].copy() + temp_memo = False + + if cls._check_shape(obj, memo): + # We update the memo every time we successfully pass a shape check + if not temp_memo: + storage.memo_stack[-1] = memo + return True + else: + return False + + def _check_shape(cls, obj, memo): + if cls.index_variadic is None: + if obj.ndim != len(cls.dims): + return False + return _check_dims(cls.dims, obj.shape, memo) + else: + if obj.ndim < len(cls.dims) - 1: + return False + i = cls.index_variadic + j = -(len(cls.dims) - i - 1) + if j == 0: + j = None + if not _check_dims(cls.dims[:i], obj.shape[:i], memo): + return False + if j is not None and not _check_dims(cls.dims[j:], obj.shape[j:], memo): + return False + variadic_dim = cls.dims[i] + if variadic_dim is not _anonymous_variadic_dim: + variadic_name = variadic_dim.name + try: + variadic_shape = memo[variadic_name] + except KeyError: + memo[variadic_name] = obj.shape[i:j] + else: + if variadic_dim.broadcastable: + new_variadic_shape = [] + obj_shape = obj.shape[i:j] + if len(variadic_shape) != len(obj_shape): + return False + for old_size, new_size in zip(variadic_shape, obj_shape): + if old_size == 1: + new_variadic_shape.append(new_size) + else: + if new_size != 1 and old_size != new_size: + return False + new_variadic_shape.append(old_size) + memo[variadic_name] = tuple(new_variadic_shape) + else: + return variadic_shape == obj.shape[i:j] + return True + + +class AbstractArray(metaclass=_MetaAbstractArray): + dtypes: List[jnp.dtype] + dims: List[_AbstractDimOrVariadicDim] + index_variadic: Optional[int] + + +class _MetaAbstractDtype(type): + def __instancecheck__(cls, obj: Any) -> NoReturn: + raise RuntimeError( + f"Do not use `isinstance(x, jaxtyping.{cls.__name__}`. If you want to " + "check just the dtype of an array, then use " + f'`jaxtyping.{cls.__name__}["..."]`.' + ) + + @ft.lru_cache(maxsize=None) + def __getitem__(cls, dim_str: str) -> _MetaAbstractArray: + if not isinstance(dim_str, str): + raise ValueError( + "Shape specification must be a string. Axes should be separated with spaces." + ) + dims = [] + index_variadic = None + for index, elem in enumerate(dim_str.split()): + if "," in elem: + # Common mistake + raise ValueError( + "Dimensions should be separated with spaces, not commas" + ) + broadcastable = False + if elem.endswith("#"): + broadcastable = True + elem = elem[:-1] + try: + elem = int(elem) + except ValueError: + if elem == "_": + elem = _anonymous_dim + elif elem == "...": + if index_variadic is not None: + raise ValueError("Cannot have multiple variadic dimensions") + index_variadic = index + elem = _anonymous_variadic_dim + elif elem[0] == "*": + if index_variadic is not None: + raise ValueError("Cannot have multiple variadic dimensions") + index_variadic = index + elem = _NamedVariadicDim(elem[1:], broadcastable) + else: + elem = _NamedDim(elem, broadcastable) + else: + elem = _FixedDim(elem, broadcastable) + dims.append(elem) + if _array_name_format == "dtype_and_shape": + name = f"{cls.__name__}['{dim_str}']" + elif _array_name_format == "array": + name = "Array" + else: + raise ValueError(f"array_name_format {_array_name_format} not recognised") + return _MetaAbstractArray( + name, + (AbstractArray,), + dict(dtypes=cls.dtypes, dims=dims, index_variadic=index_variadic), + ) + + +class AbstractDtype(metaclass=_MetaAbstractDtype): + dtypes: Union[str, List[str], Literal[_any_dtype]] + + def __init__(self, *args, **kwargs): + raise RuntimeError( + "AbstractDtype cannot be instantiated. Perhaps you wrote e.g. " + '`f32("shape")` when you mean `f32["shape"]`?' + ) + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + + dtypes = cls.dtypes + if dtypes is not _any_dtype: + if not isinstance(dtypes, list): + dtypes = [dtypes] + dtypes = [jnp.dtype(d) for d in dtypes] + cls.dtypes = dtypes + + +_bool = "bool" +_uint8 = "uint8" +_uint16 = "uint16" +_uint32 = "uint32" +_uint64 = "uint64" +_int8 = "int8" +_int16 = "int16" +_int32 = "int32" +_int64 = "int64" +_bfloat16 = "bfloat16" +_float16 = "float16" +_float32 = "float32" +_float64 = "float64" +_complex64 = "complex64" +_complex128 = "complex128" + + +def _make_dtype(_dtypes, name): + class _Cls(AbstractDtype): + dtypes = _dtypes + + _Cls.__name__ = name + _Cls.__qualname__ = name + return _Cls + + +b = _make_dtype(_bool, "b") +u8 = _make_dtype(_uint8, "u8") +u16 = _make_dtype(_uint16, "u16") +u32 = _make_dtype(_uint32, "u32") +u64 = _make_dtype(_uint64, "u64") +i8 = _make_dtype(_int8, "i8") +i16 = _make_dtype(_int16, "i16") +i32 = _make_dtype(_int32, "i32") +i64 = _make_dtype(_int64, "i64") +bf16 = _make_dtype(_bfloat16, "bf16") +f16 = _make_dtype(_float16, "f16") +f32 = _make_dtype(_float32, "f32") +f64 = _make_dtype(_float64, "f64") +c64 = _make_dtype(_complex64, "c64") +c128 = _make_dtype(_complex128, "c128") + +uints = [_uint8, _uint16, _uint32, _uint64] +ints = [_int8, _int16, _int32, _int64] +floats = [_bfloat16, _float16, _float32, _float64] +complexes = [_complex64, _complex128] + +# We match NumPy's type hierarachy in what types to provide. See the diagram at +# https://numpy.org/doc/stable/reference/arrays.scalars.html#scalars +# +# No attempt is made to match up against their character codes: all of the below are +# abstract base classes without NumPy chararacter codes. + +u = _make_dtype(uints, "u") +i = _make_dtype(ints, "i") +t = _make_dtype(uints + ints, "t") # integer +f = _make_dtype(floats, "f") +c = _make_dtype(complexes, "c") +x = _make_dtype(floats + complexes, "x") # inexact +n = _make_dtype(uints + ints + floats + complexes, "n") # number +Array = _make_dtype(_any_dtype, "Array") diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py new file mode 100644 index 0000000..5d34016 --- /dev/null +++ b/jaxtyping/decorator.py @@ -0,0 +1,38 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import functools as ft +import threading + + +storage = threading.local() +storage.memo_stack = [] + + +def jaxtyped(fn): + @ft.wraps(fn) + def wrapper(*args, **kwargs): + memo = {} + storage.memo_stack.append(memo) + try: + return fn(*args, **kwargs) + finally: + storage.memo_stack.pop() + + return wrapper diff --git a/jaxtyping/import_hook.py b/jaxtyping/import_hook.py new file mode 100644 index 0000000..a419fde --- /dev/null +++ b/jaxtyping/import_hook.py @@ -0,0 +1,272 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +# This source code is adapted from typeguard: +# https://github.com/agronholm/typeguard/blob/0dd7f7510b7c694e66a0d17d1d58d185125bad5d/src/typeguard/importhook.py +# +# Copied and adapted in compliance with the terms of typeguard's MIT license. +# The original license is reproduced here. +# +# --------- +# +# This is the MIT license: http://www.opensource.org/licenses/mit-license.php +# +# Copyright (c) Alex Grönholm +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of this +# software and associated documentation files (the "Software"), to deal in the Software +# without restriction, including without limitation the rights to use, copy, modify, merge, +# publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons +# to whom the Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or +# substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR +# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE +# FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# +# --------- + + +import ast +import sys +from importlib.abc import MetaPathFinder +from importlib.machinery import SourceFileLoader +from importlib.util import cache_from_source, decode_source +from inspect import isclass +from typing import Iterable, List, Optional, Tuple +from unittest.mock import patch + + +# The name of this function is magical +def _call_with_frames_removed(f, *args, **kwargs): + return f(*args, **kwargs) + + +def _optimized_cache_from_source(path, debug_override=None): + return cache_from_source(path, debug_override, optimization="jaxtyping") + + +class _JaxtypingTransformer(ast.NodeVisitor): + def __init__(self, *, typechecker) -> None: + self._parents: List[ast.AST] = [] + self._typechecker = typechecker + + def visit_Module(self, node: ast.Module): + # Insert "import typeguard; import jaxtping" after any "from __future__ ..." + # imports + for i, child in enumerate(node.body): + if isinstance(child, ast.ImportFrom) and child.module == "__future__": + continue + elif isinstance(child, ast.Expr) and isinstance(child.value, ast.Str): + continue # module docstring + else: + node.body.insert(i, ast.Import(names=[ast.alias("jaxtyping", None)])) + if self._typechecker is not None: + typechecker_module, _ = self._typechecker + node.body.insert( + i, ast.Import(names=[ast.alias(typechecker_module, None)]) + ) + break + + self._parents.append(node) + self.generic_visit(node) + self._parents.pop() + return node + + def visit_FunctionDef(self, node: ast.FunctionDef): + has_annotated_args = any(arg for arg in node.args.args if arg.annotation) + has_annotated_return = bool(node.returns) + if has_annotated_args or has_annotated_return: + # Place at the start of the decorator list, in case a typechecking + # annotation has been manually applied; we need to be above that. + node.decorator_list.insert( + 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 + # frequently remove annotations from functions and we'd like to + # use those annotations. + typechecker_module, typechecker_function = self._typechecker + node.decorator_list.append( + ast.Attribute( + ast.Name(id=typechecker_module, ctx=ast.Load()), + typechecker_function, + ast.Load(), + ) + ) + + self._parents.append(node) + self.generic_visit(node) + self._parents.pop() + return node + + +class _JaxtypingLoader(SourceFileLoader): + def __init__(self, *args, typechecker, **kwargs): + super().__init__(*args, **kwargs) + self._typechecker = typechecker + + def source_to_code(self, data, path, *, _optimize=-1): + source = decode_source(data) + tree = _call_with_frames_removed( + compile, + source, + path, + "exec", + ast.PyCF_ONLY_AST, + dont_inherit=True, + optimize=_optimize, + ) + tree = _JaxtypingTransformer(typechecker=self._typechecker).visit(tree) + ast.fix_missing_locations(tree) + return _call_with_frames_removed( + compile, tree, path, "exec", dont_inherit=True, optimize=_optimize + ) + + def exec_module(self, module): + # Use a custom optimization marker – the import lock should make this monkey patch safe + with patch( + "importlib._bootstrap_external.cache_from_source", + _optimized_cache_from_source, + ): + return super().exec_module(module) + + +class _JaxtypingFinder(MetaPathFinder): + """Wraps another path finder and instruments the module with `@jaxtyped` and + `@typechecked` if `should_instrument()` returns `True`. + + Should not be used directly, but rather via `install_import_hook`. + """ + + def __init__(self, modules, original_pathfinder, typechecker): + self.modules = modules + self._original_pathfinder = original_pathfinder + self._typechecker = typechecker + + def find_spec(self, fullname, path=None, target=None): + if self.should_instrument(fullname): + spec = self._original_pathfinder.find_spec(fullname, path, target) + if spec is not None and isinstance(spec.loader, SourceFileLoader): + spec.loader = _JaxtypingLoader( + spec.loader.name, spec.loader.path, typechecker=self._typechecker + ) + return spec + + return None + + def should_instrument(self, module_name: str) -> bool: + """Determine whether the module with the given name should be instrumented. + + **Arguments:** + + - `module_name`: the full name of the module that is about to be imported + (e.g. ``xyz.abc``) + """ + for module in self.modules: + if module_name == module or module_name.startswith(module + "."): + return True + + return False + + +class ImportHookManager: + def __init__(self, hook: MetaPathFinder): + self.hook = hook + + def __enter__(self): + pass + + def __exit__(self, exc_type, exc_val, exc_tb): + self.uninstall() + + def uninstall(self): + try: + sys.meta_path.remove(self.hook) + except ValueError: + pass # already removed + + +# Deliberately no default for `typechecker` so that folks must opt-in to not having +# a typechecker. +def install_import_hook( + modules: Iterable[str], typechecker: Optional[Tuple[str, str]] +) -> ImportHookManager: + """Automatically apply `@jaxtyped`, and optionally a type checker, to all classes and functions. + + It will only be applied to modules loaded **after** this hook has been installed. + + **Arguments:**: + + - `packages`: the names of the modules in which to automatically apply `@jaxtyped` + and `@typechecked`. + - `typechecker`: the module and function of the typechecker you want to use, as a + 2-tuple of strings. For example `typechecker=("typeguard", "typechecked")` or + `typechecker=("beartype", "beartype")`. You may pass `typechecker=None` if you + do not want to automatically decorate with a typechecker as well; e.g. if you + have a codebase that already has these decorators. + + **Returns:** + + A context manager that uninstalls the hook on exit, or when you call `.uninstall()`. + + **Example:** + + Typically you should apply this import hook at the entry point for your own scripts: + ```python + # entry_point.py + from jaxtyped import install_import_hook + install_import_hook("main", ("beartype", "beartype")) + import main + ... # do whatever you're doing + + # main.py + from jaxtyped import f32 + + def f(x: f32["b c"]): + pass + ``` + Which as you can see means you never to import `@jaxtyped`, nor do you need to + import the typechecker directly (e.g. `beartype.beartype` or + `typeguard.typechecked`). + """ + if isinstance(modules, str): + modules = [modules] + + for i, finder in enumerate(sys.meta_path): + if ( + isclass(finder) + and finder.__name__ == "PathFinder" + and hasattr(finder, "find_spec") + ): + break + else: + raise RuntimeError("Cannot find a PathFinder in sys.meta_path") + + hook = _JaxtypingFinder(modules, finder, typechecker) + sys.meta_path.insert(0, hook) + return ImportHookManager(hook) diff --git a/jaxtyping/pytest_plugin.py b/jaxtyping/pytest_plugin.py new file mode 100644 index 0000000..dafd6aa --- /dev/null +++ b/jaxtyping/pytest_plugin.py @@ -0,0 +1,55 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import sys + +from .import_hook import install_import_hook + + +def pytest_addoption(parser): + group = parser.getgroup("jaxtyping") + group.addoption( + "--jaxtyping-packages", + action="store", + help="comma separated name list of packages and modules to instrument for " + "type checking with jaxtyping. The last element in the list should be the " + "type checker to use, e.g. " + "--jaxtyping-packages=foopackage,barpackage,typeguard.typechecked", + ) + + +def pytest_configure(config): + value = config.getoption("jaxtyping_packages") + if not value: + return + + packages = [pkg.strip() for pkg in value.split(",")] + *packages, typechecker = packages + + already_imported_packages = sorted( + package for package in packages if package in sys.modules + ) + if already_imported_packages: + message = ( + "jaxtyping cannot check these packages because they " + "are already imported: {}" + ) + raise RuntimeError(message.format(", ".join(already_imported_packages))) + + install_import_hook(packages, typechecker.rsplit(".", 1)) diff --git a/jaxtyping/pytree_type.py b/jaxtyping/pytree_type.py new file mode 100644 index 0000000..7a7d894 --- /dev/null +++ b/jaxtyping/pytree_type.py @@ -0,0 +1,65 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import functools as ft + +import jax +import typeguard + + +class _MetaPyTree(type): + def __call__(self, *args, **kwargs): + raise RuntimeError("PyTree cannot be instantiated") + + def __instancecheck__(cls, obj): + return True + + @ft.lru_cache(maxsize=None) + def __getitem__(cls, item): + return _MetaSubscriptPyTree(f"PyTree[{item.__name__}]", (), {"leaftype": item}) + + +class _MetaSubscriptPyTree(type): + def __call__(self, *args, **kwargs): + raise RuntimeError("PyTree cannot be instantiated") + + def __instancecheck__(cls, obj): + # We could use `isinstance` here but that would fail for more complicated + # types, e.g. PyTree[Tuple[int]]. So at least internally we make a particular + # choice of typechecker. + # + # Deliberately not using @jaxtyped so that we share the same `memo` as whatever + # dynamic context we're currently in. + @typeguard.typechecked + def accepts_leaftype(x: cls.leaftype): + pass + + def is_leaftype(x): + try: + accepts_leaftype(x) + except TypeError: + return False + else: + return True + + leaves = jax.tree_leaves(obj, is_leaf=is_leaftype) + return all(map(is_leaftype, leaves)) + + +PyTree = _MetaPyTree("PyTree", (), {}) diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..9776fcd --- /dev/null +++ b/setup.py @@ -0,0 +1,88 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import pathlib +import re + +import setuptools + + +_here = pathlib.Path(__file__).resolve().parent + + +name = "jaxtyping" + +# for simplicity we actually store the version in the __version__ attribute in the source +with open(_here / name / "__init__.py") as f: + meta_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", f.read(), re.M) + if meta_match: + version = meta_match.group(1) + else: + raise RuntimeError("Unable to find __version__ string.") + +author = "Patrick Kidger" + +author_email = "contact@kidger.site" + +description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees." + +with open(_here / "README.md", "r") as f: + readme = f.read() + +url = "https://github.com/google/jaxtyping" + +license = "MIT" + +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Natural Language :: English", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Scientific/Engineering :: Mathematics", +] + +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"] + +entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"]) + +setuptools.setup( + name=name, + version=version, + author=author, + author_email=author_email, + maintainer=author, + maintainer_email=author_email, + description=description, + long_description=readme, + long_description_content_type="text/markdown", + url=url, + license=license, + classifiers=classifiers, + zip_safe=False, + python_requires=python_requires, + install_requires=install_requires, + entry_points=entry_points, + packages=[name], +) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..f549b6b --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,42 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import random + +import beartype +import jax.random as jr +import pytest +import typeguard + + +@pytest.fixture(params=[typeguard.typechecked, beartype.beartype]) +def typecheck(request): + return request.param + + +@pytest.fixture() +def getkey(): + def _getkey(): + # Not sure what the maximum actually is but this will do + return jr.PRNGKey(random.randint(0, 2**31 - 1)) + + return _getkey + + +ParamException = (TypeError, beartype.roar.BeartypeCallHintParamViolation) diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..8e5289e --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,30 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import beartype +import equinox as eqx + + +ParamError = (TypeError, beartype.roar.BeartypeCallHintParamViolation) +ReturnError = (TypeError, beartype.roar.BeartypeCallHintReturnViolation) + + +@eqx.filter_jit +def make_mlp(key): + return eqx.nn.MLP(2, 2, 2, 2, key=key) diff --git a/tests/import_hook_tester_beartype.py b/tests/import_hook_tester_beartype.py new file mode 100644 index 0000000..6c5d01d --- /dev/null +++ b/tests/import_hook_tester_beartype.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import jax.numpy as jnp +import pytest +from helpers import ParamError + +from jaxtyping import f32 + + +def g(x: f32[" b"]): + pass + + +g(jnp.array([1.0])) +with pytest.raises(ParamError): + g(jnp.array(1)) diff --git a/tests/import_hook_tester_broken_checker.py b/tests/import_hook_tester_broken_checker.py new file mode 100644 index 0000000..6c5d01d --- /dev/null +++ b/tests/import_hook_tester_broken_checker.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import jax.numpy as jnp +import pytest +from helpers import ParamError + +from jaxtyping import f32 + + +def g(x: f32[" b"]): + pass + + +g(jnp.array([1.0])) +with pytest.raises(ParamError): + g(jnp.array(1)) diff --git a/tests/import_hook_tester_transitive/__init__.py b/tests/import_hook_tester_transitive/__init__.py new file mode 100644 index 0000000..1bb4722 --- /dev/null +++ b/tests/import_hook_tester_transitive/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +from . import another_file diff --git a/tests/import_hook_tester_transitive/another_file.py b/tests/import_hook_tester_transitive/another_file.py new file mode 100644 index 0000000..6c5d01d --- /dev/null +++ b/tests/import_hook_tester_transitive/another_file.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import jax.numpy as jnp +import pytest +from helpers import ParamError + +from jaxtyping import f32 + + +def g(x: f32[" b"]): + pass + + +g(jnp.array([1.0])) +with pytest.raises(ParamError): + g(jnp.array(1)) diff --git a/tests/import_hook_tester_typeguard.py b/tests/import_hook_tester_typeguard.py new file mode 100644 index 0000000..6c5d01d --- /dev/null +++ b/tests/import_hook_tester_typeguard.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import jax.numpy as jnp +import pytest +from helpers import ParamError + +from jaxtyping import f32 + + +def g(x: f32[" b"]): + pass + + +g(jnp.array([1.0])) +with pytest.raises(ParamError): + g(jnp.array(1)) diff --git a/tests/requirements.txt b/tests/requirements.txt new file mode 100644 index 0000000..dd54a1e --- /dev/null +++ b/tests/requirements.txt @@ -0,0 +1,4 @@ +equinox>=0.5.3 +pytest>=7.0.1 +beartype>=0.10.4 +typeguard>=2.13.3 diff --git a/tests/test_array.py b/tests/test_array.py new file mode 100644 index 0000000..96fdb9b --- /dev/null +++ b/tests/test_array.py @@ -0,0 +1,331 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import jax.numpy as jnp +import jax.random as jr +import pytest +from helpers import ParamError, ReturnError + +from jaxtyping import Array, f, f32, jaxtyped + + +def test_basic(typecheck): + @jaxtyped + @typecheck + def g(x: Array["..."]): + pass + + g(jnp.array(1.0)) + + +def test_return(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f["b c"]) -> f["c b"]: + return jnp.transpose(x) + + g(jr.normal(getkey(), (3, 4))) + + @jaxtyped + @typecheck + def h(x: f["b c"]) -> f["b c"]: + return jnp.transpose(x) + + with pytest.raises(ReturnError): + h(jr.normal(getkey(), (3, 4))) + + +def test_two_args(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: Array["b c"], y: Array["c d"]): + return x @ y + + g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (4, 5))) + with pytest.raises(ParamError): + g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4))) + + @jaxtyped + @typecheck + def h(x: Array["b c"], y: Array["c d"]) -> Array["b d"]: + return x @ y + + h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (4, 5))) + with pytest.raises(ParamError): + h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4))) + + +def test_any_dtype(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: Array["a b"]) -> Array["a b"]: + return x + + g(jr.normal(getkey(), (3, 4))) + g(jnp.array([[True, False]])) + g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int8)) + g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint16)) + g(jr.normal(getkey(), (3, 4), dtype=jnp.complex128)) + g(jr.normal(getkey(), (3, 4), dtype=jnp.bfloat16)) + + with pytest.raises(ParamError): + g(jr.normal(getkey(), (1,))) + + +def test_nested_jaxtyped(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32["b c"], transpose: bool) -> f32["c b"]: + return h(x, transpose) + + @jaxtyped + @typecheck + def h(x: f32["c b"], transpose: bool) -> f32["b c"]: + if transpose: + return jnp.transpose(x) + else: + return x + + g(jr.normal(getkey(), (2, 3)), True) + g(jr.normal(getkey(), (3, 3)), True) + g(jr.normal(getkey(), (3, 3)), False) + with pytest.raises(ReturnError): + g(jr.normal(getkey(), (2, 3)), False) + + +def test_nested_nojaxtyped(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32["b c"]): + return h(x) + + @typecheck + def h(x: f32["c b"]): + return x + + with pytest.raises(ParamError): + g(jr.normal(getkey(), (2, 3))) + + +def test_isinstance(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32["b c"]) -> f32[" z"]: + y = jnp.transpose(x) + assert isinstance(y, f32["c b"]) + assert not isinstance( + y, f32["b z"] + ) # z left unbound as b!=c (unless x symmetric, which it isn't) + out = jr.normal(getkey(), (500,)) + assert isinstance(out, f32["z"]) # z now bound + return out + + g(jr.normal(getkey(), (2, 3))) + + +def test_fixed(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32["4 5 foo"], y: f32[" foo"]) -> f32["4 5"]: + return x @ y + + a = jr.normal(getkey(), (4, 5, 2)) + b = jr.normal(getkey(), (2,)) + assert g(a, b).shape == (4, 5) + + c = jr.normal(getkey(), (3, 5, 2)) + with pytest.raises(ParamError): + g(c, b) + + +def test_anonymous(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32["foo _"], y: f32[" _"]): + pass + + a = jr.normal(getkey(), (3, 4)) + b = jr.normal(getkey(), (5,)) + g(a, b) + + +def test_named_variadic(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32["*batch foo"], y: f32[" *batch"], z: f32[" foo"]): + pass + + c = jr.normal(getkey(), (5,)) + + a1 = jr.normal(getkey(), (5,)) + b1 = jr.normal(getkey(), ()) + g(a1, b1, c) + + a2 = jr.normal(getkey(), (3, 5)) + b2 = jr.normal(getkey(), (3,)) + g(a2, b2, c) + + with pytest.raises(ParamError): + g(a1, b2, c) + with pytest.raises(ParamError): + g(a2, b1, c) + + @jaxtyped + @typecheck + def h(x: f32[" foo *batch"], y: f32[" foo *batch bar"]): + pass + + a = jr.normal(getkey(), (4,)) + b = jr.normal(getkey(), (4, 3)) + c = jr.normal(getkey(), (3, 4)) + h(a, b) + with pytest.raises(ParamError): + h(a, c) + with pytest.raises(ParamError): + h(b, c) + + +def test_anonymous_variadic(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32["... foo"], y: f32[" foo"]): + pass + + a1 = jr.normal(getkey(), (5,)) + a2 = jr.normal(getkey(), (3, 5)) + a3 = jr.normal(getkey(), (3, 4, 5)) + b = jr.normal(getkey(), (5,)) + c = jr.normal(getkey(), (1,)) + g(a1, b) + g(a2, b) + g(a3, b) + with pytest.raises(ParamError): + g(a1, c) + with pytest.raises(ParamError): + g(a2, c) + with pytest.raises(ParamError): + g(a3, c) + + +def test_broadcast_fixed(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32["4#"]): + pass + + g(jr.normal(getkey(), (4,))) + g(jr.normal(getkey(), (1,))) + + with pytest.raises(ParamError): + g(jr.normal(getkey(), (3,))) + + +def test_broadcast_named(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32[" foo#"], y: f32[" foo#"]): + pass + + a = jr.normal(getkey(), (3,)) + b = jr.normal(getkey(), (4,)) + c = jr.normal(getkey(), (1,)) + + g(a, a) + g(b, b) + g(c, c) + g(a, c) + g(b, c) + g(c, a) + g(c, b) + + with pytest.raises(ParamError): + g(a, b) + with pytest.raises(ParamError): + g(b, a) + + +def test_broadcast_variadic_named(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: f32[" *foo#"], y: f32[" *foo#"]): + pass + + a = jr.normal(getkey(), (3,)) + b = jr.normal(getkey(), (4,)) + c = jr.normal(getkey(), (4, 4)) + d = jr.normal(getkey(), (5, 6)) + + j = jr.normal(getkey(), (1,)) + k = jr.normal(getkey(), (1, 4)) + l = jr.normal(getkey(), (5, 1)) # noqa: E741 + m = jr.normal(getkey(), (1, 1)) + n = jr.normal(getkey(), (2, 1)) + o = jr.normal(getkey(), (1, 6)) + + g(a, a) + g(b, b) + g(c, c) + g(d, d) + with pytest.raises(ParamError): + g(a, b) + with pytest.raises(ParamError): + g(a, c) + with pytest.raises(ParamError): + g(b, c) + with pytest.raises(ParamError): + g(a, b) + with pytest.raises(ParamError): + g(d, b) + + g(a, j) + g(b, j) + with pytest.raises(ParamError): + g(c, j) + with pytest.raises(ParamError): + g(d, j) + with pytest.raises(ParamError): + g(b, k) + g(c, k) + with pytest.raises(ParamError): + g(d, k) + with pytest.raises(ParamError): + g(c, l) + g(d, l) + with pytest.raises(ParamError): + g(a, m) + g(c, m) + g(d, m) + with pytest.raises(ParamError): + g(a, n) + with pytest.raises(ParamError): + g(b, n) + with pytest.raises(ParamError): + g(c, n) + with pytest.raises(ParamError): + g(d, n) + g(o, d) + with pytest.raises(ParamError): + g(o, c) + with pytest.raises(ParamError): + g(o, a) + + +def test_no_commas(typecheck, getkey): + with pytest.raises(ValueError): + f32["foo, bar"] diff --git a/tests/test_import_hook.py b/tests/test_import_hook.py new file mode 100644 index 0000000..2881394 --- /dev/null +++ b/tests/test_import_hook.py @@ -0,0 +1,56 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +import pytest + +from jaxtyping import install_import_hook + + +def test_import_hook_typeguard(): + hook = install_import_hook( + "import_hook_tester_typeguard", ("typeguard", "typechecked") + ) + import import_hook_tester_typeguard # noqa: F401 + + hook.uninstall() + + +def test_import_hook_beartype(): + hook = install_import_hook("import_hook_tester_beartype", ("beartype", "beartype")) + import import_hook_tester_beartype # noqa: F401 + + hook.uninstall() + + +def test_import_hook_transitive(): + hook = install_import_hook( + "import_hook_tester_transitive", ("typeguard", "typechecked") + ) + import import_hook_tester_transitive # noqa: F401 + + hook.uninstall() + + +def test_import_hook_broken_checker(): + hook = install_import_hook( + "import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist") + ) + with pytest.raises(AttributeError): + import import_hook_tester_broken_checker # noqa: F401 + hook.uninstall() diff --git a/tests/test_pytree.py b/tests/test_pytree.py new file mode 100644 index 0000000..f549cb4 --- /dev/null +++ b/tests/test_pytree.py @@ -0,0 +1,156 @@ +# Copyright (c) 2022 Google LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of +# this software and associated documentation files (the "Software"), to deal in +# the Software without restriction, including without limitation the rights to +# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +# the Software, and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +from typing import Tuple, Union + +import equinox as eqx +import jax +import jax.numpy as jnp +import jax.random as jr +import pytest +from helpers import make_mlp, ParamError + +from jaxtyping import f, jaxtyped, PyTree + + +def test_direct(typecheck): + @typecheck + def g(x: PyTree): + pass + + g(1) + g({"a": jnp.array(1), "b": [object()]}) + g(object()) + + @typecheck + def h() -> PyTree: + return object() + + h() + + +def test_subscript(getkey, typecheck): + @typecheck + def g(x: PyTree[int]): + pass + + g(1) + g([1, 2, {"a": 3}]) + g(jax.tree_map(lambda _: 1, make_mlp(getkey()))) + + with pytest.raises(ParamError): + g(object()) + with pytest.raises(ParamError): + g("hi") + with pytest.raises(ParamError): + g([1, 2, {"a": 3}, "bye"]) + + +def test_leaf_pytrees(getkey, typecheck): + @typecheck + def g(x: PyTree[eqx.nn.MLP]): + pass + + g(make_mlp(getkey())) + g([make_mlp(getkey()), make_mlp(getkey()), {"a": make_mlp(getkey())}]) + + with pytest.raises(ParamError): + g([1, 2]) + with pytest.raises(ParamError): + g([1, 2, make_mlp()]) + + +def test_nested_pytrees(getkey, typecheck): + # PyTree[...] is logically equivalent to PyTree[PyTree[...]] + @typecheck + def g(x: PyTree[PyTree[eqx.nn.MLP]]): + pass + + g(make_mlp(getkey())) + g([make_mlp(getkey()), make_mlp(getkey()), {"a": make_mlp(getkey())}]) + + with pytest.raises(ParamError): + g([1, 2]) + with pytest.raises(ParamError): + g([1, 2, make_mlp()]) + + +def test_pytree_array(typecheck): + @jaxtyped + @typecheck + def g(x: PyTree[f["..."]]): + pass + + g(jnp.array(1.0)) + g([jnp.array(1.0), jnp.array(1.0), {"a": jnp.array(1.0)}]) + with pytest.raises(ParamError): + g(jnp.array(1)) + with pytest.raises(ParamError): + g(1.0) + + +def test_pytree_shaped_array(typecheck, getkey): + @jaxtyped + @typecheck + def g(x: PyTree[f["b c"]]): + pass + + g(jnp.array([[1.0]])) + g([jr.normal(getkey(), (1, 1)), jr.normal(getkey(), (1, 1))]) + g([jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (3, 4))]) + with pytest.raises(ParamError): + g(jnp.array(1.0)) + with pytest.raises(ParamError): + g(jnp.array([[1]])) + with pytest.raises(ParamError): + g([jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (1, 1))]) + with pytest.raises(ParamError): + g([jr.normal(getkey(), (1, 1)), jr.normal(getkey(), (3, 4))]) + + +def test_pytree_union(typecheck): + @typecheck + def g(x: PyTree[Union[int, str]]): + pass + + g([1]) + g(["hi"]) + g([1, "hi"]) + with pytest.raises(ParamError): + g(object()) + + +def test_pytree_tuple(typecheck): + @typecheck + def g(x: PyTree[Tuple[int, int]]): + pass + + g((1, 1)) + g([(1, 1)]) + g([(1, 1), {"a": (1, 1)}]) + with pytest.raises(ParamError): + g(object()) + with pytest.raises(ParamError): + g(1) + with pytest.raises(ParamError): + g((1, 2, 3)) + with pytest.raises(ParamError): + g([1, 1]) + with pytest.raises(ParamError): + g([(1, 1), "hi"])