wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-09 11:24:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

initial commit

Browse Source
This commit is contained in:
Patrick Kidger
2022-07-01 21:01:46 +00:00
commit 7ac6ee04a8
30 changed files with 2207 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.flake8
+4
View File
@@ -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
.github/workflows/release.yml
+45
View File
@@ -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 }}
.github/workflows/run_tests.yml
+53
View File
@@ -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
.gitignore
+3
View File
@@ -0,0 +1,3 @@
**/__pycache__
*.egg-info
.isort.cfg
+5
View File
@@ -0,0 +1,5 @@
[settings]
force_alphabetical_sort_within_sections=true
lines_after_imports=2
profile=black
treat_comments_as_code=true
.pre-commit-config.yaml
+38
View File
@@ -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
API.md
+196
View File
@@ -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 <ins>n</ins>umber)
- Any floating or complex: `x` (for ine<ins>x</ins>act)
- 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 in<ins>t</ins>eger)
- 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"]`.
CONTRIBUTING.md
+54
View File
@@ -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
<https://cla.developers.google.com/> 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.
FAQ.md
+43
View File
@@ -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.
LICENSE
+51
View File
@@ -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.
MANIFEST.in
+2
View File
@@ -0,0 +1,2 @@
include LICENSE
prune tests
README.md
+56
View File
@@ -0,0 +1,56 @@
<h1 align="center">jaxtyping</h1>
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.
jaxtyping/__init__.py
+54
View File
@@ -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"
jaxtyping/array_types.py
+314
View File
@@ -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")
jaxtyping/decorator.py
+38
View File
@@ -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
jaxtyping/import_hook.py
+272
View File
@@ -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)
jaxtyping/pytest_plugin.py
+55
View File
@@ -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))
jaxtyping/pytree_type.py
+65
View File
@@ -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", (), {})
setup.py
+88
View File
@@ -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],
)
tests/conftest.py
+42
View File
@@ -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)
tests/helpers.py
+30
View File
@@ -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)
tests/import_hook_tester_beartype.py
+33
View File
@@ -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))
tests/import_hook_tester_broken_checker.py
+33
View File
@@ -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))
tests/import_hook_tester_transitive/__init__.py
+20
View File
@@ -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
tests/import_hook_tester_transitive/another_file.py
+33
View File
@@ -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))
tests/import_hook_tester_typeguard.py
+33
View File
@@ -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))
tests/requirements.txt
+4
View File
@@ -0,0 +1,4 @@
equinox>=0.5.3
pytest>=7.0.1
beartype>=0.10.4
typeguard>=2.13.3
tests/test_array.py
+331
View File
@@ -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"]
tests/test_import_hook.py
+56
View File
@@ -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()
tests/test_pytree.py
+156
View File
@@ -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"])