Static typing fixes

Also added support for matching dtypes using regexes, and for nesting e.g. Shaped[Float[Array, "dim1 dim2"], "dim3"].

.flake8

@@ -1,4 +0,0 @@
-[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/build_docs.yml

@@ -0,0 +1,39 @@
+name: Build docs
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        python-version: [ 3.11 ]
+        os: [ ubuntu-latest ]
+    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 .
+          python -m pip install -r docs/requirements.txt
+
+      - name: Build docs
+        run: |
+          mkdocs build
+          mkdocs build  # twice, see https://github.com/patrick-kidger/pytkdocs_tweaks
+
+      - name: Upload docs
+        uses: actions/upload-artifact@v2
+        with:
+          name: docs
+          path: site  # where `mkdocs build` puts the built site

.github/workflows/release.yml

@@ -29,11 +29,12 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Release
-        uses: patrick-kidger/action_update_python_project@v1
+        uses: patrick-kidger/action_update_python_project@v2
         with:
-            python-version: "3.8"
+            python-version: "3.11"
             test-script: |
-                python -m pip install pytest beartype equinox jaxlib
+                python -m pip install pytest beartype equinox jaxlib cloudpickle
+                python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
                 cp -r ${{ github.workspace }}/test ./test
                 pytest
             pypi-token: ${{ secrets.pypi_token }}

.github/workflows/run_tests.yml

@@ -26,7 +26,7 @@ jobs:
   run-tests:
     strategy:
       matrix:
-        python-version: [ 3.7, 3.8, 3.9 ]
+        python-version: [ 3.11 ]
         os: [ ubuntu-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}
@@ -42,7 +42,8 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest wheel beartype equinox jaxlib
+          python -m pip install pytest wheel beartype equinox jaxlib cloudpickle
+          python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
 
       - name: Checks with pre-commit
         uses: pre-commit/action@v2.0.3

.gitignore

@@ -2,4 +2,5 @@
 *.egg-info
 build/
 dist/
-
+site/
+.all_objects.cache

.isort.cfg

@@ -1,6 +0,0 @@
-[settings]
-force_alphabetical_sort_within_sections=true
-lines_after_imports=2
-profile=black
-treat_comments_as_code=true
-extra_standard_library=typing_extensions

.pre-commit-config.yaml

@@ -22,17 +22,7 @@ repos:
     rev: 22.3.0
     hooks:
     - id: black
-  - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.2.3
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: 'v0.0.255'
     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
+      - id: ruff

API.md

@@ -1,196 +0,0 @@
-# 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

@@ -28,7 +28,8 @@ Now make your changes. Make sure to include additional tests if necessary.
 Next verify the tests all pass:
 
 ```bash
-pip install pytest
+pip install pytest cloudpickle
+pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
 pytest
 ```
 

FAQ.md

@@ -1,43 +0,0 @@
-# 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.

MANIFEST.in

@@ -1,2 +0,0 @@
-include LICENSE
-prune tests

README.md

@@ -1,21 +1,25 @@
 <h1 align="center">jaxtyping</h1>
 
-Type annotations **and runtime checking** for:
+Type annotations **and runtime type-checking** for:
 
-1. shape and dtype of [JAX](https://github.com/google/jax) arrays;
+1. shape and dtype of [JAX](https://github.com/google/jax) arrays; *(Now also supports PyTorch, NumPy, and TensorFlow!)*
 2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
 
+
 **For example:**
 ```python
-from jaxtyping import f32, PyTree
+from jaxtyping import Array, Float, PyTree
 
-def matrix_multiply(x: f32["dim1 dim2"], y: f32["dim2 dim3"]) -> f32["dim1 dim3"]:
+# Accepts floating-point 2D arrays with matching dimensions
+def matrix_multiply(x: Float[Array, "dim1 dim2"],
+                    y: Float[Array, "dim2 dim3"]
+                  ) -> Float[Array, "dim1 dim3"]:
     ...
 
 def accepts_pytree_of_ints(x: PyTree[int]):
     ...
 
-def accepts_pytree_of_arrays(x: PyTree[f32["batch c1 c2"]]):
+def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
     ...
 ```
 
@@ -25,15 +29,15 @@ def accepts_pytree_of_arrays(x: PyTree[f32["batch c1 c2"]]):
 pip install jaxtyping
 ```
 
-Requires JAX 0.3.4+.
+Requires Python 3.8+.
 
-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).
+JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/etc.
+
+The annotations provided by jaxtyping are compatible with runtime type-checking packages, so it is common to also install one of these. 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)
+Available at [https://docs.kidger.site/jaxtyping](https://docs.kidger.site/jaxtyping).
 
 ## Finally
 
@@ -43,14 +47,10 @@ Neural networks: [Equinox](https://github.com/patrick-kidger/equinox).
 
 Numerical differential equation solvers: [Diffrax](https://github.com/patrick-kidger/diffrax).
 
+Computer vision models: [Eqxvision](https://github.com/paganpasta/eqxvision).
+
 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.

docs/.citation.md

@@ -0,0 +1,12 @@
+If you found this library to be useful in academic work, then please cite: ([arXiv link](https://arxiv.org/abs/2111.00254))
+
+```bibtex
+@article{kidger2021equinox,
+    author={Patrick Kidger and Cristian Garcia},
+    title={{E}quinox: neural networks in {JAX} via callable {P}y{T}rees and filtered transformations},
+    year={2021},
+    journal={Differentiable Programming workshop at Neural Information Processing Systems 2021}
+}
+```
+
+(Also consider starring the project [on GitHub](https://github.com/patrick-kidger/equinox).)

docs/.htaccess

@@ -0,0 +1 @@
+ErrorDocument 404 /jaxtyping/404.html

docs/_overrides/partials/source.html

@@ -0,0 +1,20 @@
+{% import "partials/language.html" as lang with context %}
+<a href="{{ config.repo_url }}" title="{{ lang.t('source.link.title') }}" class="md-source" data-md-component="source">
+  <div class="md-source__icon md-icon">
+    {% set icon = config.theme.icon.repo or "fontawesome/brands/git-alt" %}
+    {% include ".icons/" ~ icon ~ ".svg" %}
+  </div>
+  <div class="md-source__repository">
+    {{ config.repo_name }}
+  </div>
+</a>
+{% if config.theme.twitter_url %}
+<a href="{{ config.theme.twitter_url }}" title="Go to Twitter" class="md-source">
+  <div class="md-source__icon md-icon">
+    {% include ".icons/fontawesome/brands/twitter.svg" %}
+  </div>
+  <div class="md-source__repository">
+    {{ config.theme.twitter_name }}
+  </div>
+</a>
+{% endif %}

docs/_static/.README.md

@@ -0,0 +1,2 @@
+The favicon is `math-integral` from https://materialdesignicons.com, found by way of https://pictogrammers.com.
+(The logo is `math-integral-box`.)

docs/_static/custom_css.css

@@ -0,0 +1,167 @@
+/* Fix /page#foo going to the top of the viewport and being hidden by the navbar */
+html {
+    scroll-padding-top: 50px;
+}
+
+/* Fit the Twitter handle alongside the GitHub one in the top right. */
+
+div.md-header__source {
+    width: revert;
+    max-width: revert;
+}
+
+a.md-source {
+    display: inline-block;
+}
+
+.md-source__repository {
+    max-width: 100%;
+}
+
+/* Emphasise sections of nav on left hand side */
+
+nav.md-nav {
+  padding-left: 5px;
+}
+
+nav.md-nav--secondary {
+    border-left: revert !important;
+}
+
+.md-nav__title {
+  font-size: 0.9rem;
+}
+
+.md-nav__item--section > .md-nav__link {
+  font-size: 0.9rem;
+}
+
+/* Indent autogenerated documentation */
+
+div.doc-contents {
+  padding-left: 25px;
+  border-left: 4px solid rgba(230, 230, 230);
+}
+
+/* Increase visibility of splitters "---" */
+
+[data-md-color-scheme="default"] .md-typeset hr {
+    border-bottom-color: rgb(0, 0, 0);
+    border-bottom-width: 1pt;
+}
+
+[data-md-color-scheme="slate"] .md-typeset hr {
+    border-bottom-color: rgb(230, 230, 230);
+}
+
+/* More space at the bottom of the page */
+
+.md-main__inner {
+  margin-bottom: 1.5rem;
+}
+
+/* Remove prev/next footer buttons */
+
+.md-footer__inner {
+    display: none;
+}
+
+/* Change font sizes */
+
+html {
+    /* Decrease font size for overall webpage
+       Down from 137.5% which is the Material default */
+    font-size: 110%;
+}
+
+.md-typeset .admonition {
+    /* Increase font size in admonitions */
+    font-size: 100% !important;
+}
+
+.md-typeset details {
+    /* Increase font size in details */
+    font-size: 100% !important;
+}
+
+.md-typeset h1 {
+    font-size: 1.6rem;
+}
+
+.md-typeset h2 {
+    font-size: 1.5rem;
+}
+
+.md-typeset h3 {
+    font-size: 1.3rem;
+}
+
+.md-typeset h4 {
+    font-size: 1.1rem;
+}
+
+.md-typeset h5 {
+    font-size: 0.9rem;
+}
+
+.md-typeset h6 {
+    font-size: 0.8rem;
+}
+
+/* Bugfix: remove the superfluous parts generated when doing:
+
+??? Blah
+
+    ::: library.something
+*/
+
+.md-typeset details .mkdocstrings > h4 {
+    display: none;
+}
+
+.md-typeset details .mkdocstrings > h5 {
+    display: none;
+}
+
+/* Change default colours for <a> tags */
+
+[data-md-color-scheme="default"] {
+    --md-typeset-a-color: rgb(0, 189, 164) !important;
+}
+[data-md-color-scheme="slate"] {
+    --md-typeset-a-color: rgb(0, 189, 164) !important;
+}
+
+/* Highlight functions, classes etc. type signatures. Really helps to make clear where
+   one item ends and another begins. */
+
+[data-md-color-scheme="default"] {
+    --doc-heading-color: #DDD;
+    --doc-heading-border-color: #CCC;
+    --doc-heading-color-alt: #F0F0F0;
+}
+[data-md-color-scheme="slate"] {
+    --doc-heading-color: rgb(25,25,33);
+    --doc-heading-border-color: rgb(25,25,33);
+    --doc-heading-color-alt: rgb(33,33,44);
+    --md-code-bg-color: rgb(38,38,50);
+}
+
+h4.doc-heading {
+    /* NOT var(--md-code-bg-color) as that's not visually distinct from other code blocks.*/
+    background-color: var(--doc-heading-color);
+    border: solid var(--doc-heading-border-color);
+    border-width: 1.5pt;
+    border-radius: 2pt;
+    padding: 0pt 5pt 2pt 5pt;
+}
+h5.doc-heading, h6.heading {
+    background-color: var(--doc-heading-color-alt);
+    border-radius: 2pt;
+    padding: 0pt 5pt 2pt 5pt;
+}
+
+/* Make errors in notebooks have scrolling */
+.output_error > pre {
+    overflow: auto;
+}

docs/_static/mathjax.js

@@ -0,0 +1,16 @@
+window.MathJax = {
+  tex: {
+    inlineMath: [["\\(", "\\)"]],
+    displayMath: [["\\[", "\\]"]],
+    processEscapes: true,
+    processEnvironments: true
+  },
+  options: {
+    ignoreHtmlClass: ".*|",
+    processHtmlClass: "arithmatex"
+  }
+};
+
+document$.subscribe(() => { 
+  MathJax.typesetPromise()
+})

docs/api/advanced-features.md

@@ -0,0 +1,13 @@
+# Advanced features
+
+## Abstract base classes
+
+::: jaxtyping.AbstractDtype
+    selection:
+        members:
+            false
+
+::: jaxtyping.AbstractArray
+    selection:
+        members:
+            false

docs/api/array.md

@@ -0,0 +1,112 @@
+# Array annotations
+
+The shape and dtypes of arrays can be annotated in the form `dtype[array, shape]`, such as `Float[Array, "batch channels"]`.
+
+## Shape
+
+**Symbols**
+
+The shape should be a string of space-separated symbols, such as `"a b c d"`. Each symbol can be either an:
+
+- `int`: fixed-size axis, e.g. `"28 28"`.
+- `str`: variable-size axis, e.g. `"channels"`.
+- A symbolic expression (without spaces!) in terms of other variable-size axes, e.g.  
+    `def remove_last(x: Float[Array, "dim"]) -> Float[Array, "dim-1"]`.
+
+When calling a function, variable-size axes and symbolic axes will be matched up across all arguments and checked for consistency. (See [Runtime type checking](./runtime-type-checking.md).)
+
+**Modifiers**
+
+In addition some modifiers can be applied:
+
+- Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `"*batch c h w"` will match zero or more batch axes.
+- Prepend `#` to a dimension to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g.  
+    `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`.
+- Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`.
+- Documentation-only names (i.e. they're ignored by jaxtyping) can be handled by prepending a name followed by `=` e.g. `Float[Array, "rows=4 cols=3"]`.
+
+When using multiple modifiers, their order does not matter.
+
+As a special case:
+
+- `...`: anonymous zero or more axes (equivalent to `*_`) e.g. `"... c h w"`
+
+**Notes**
+
+- To denote a scalar shape use `""`, e.g. `Float[Array, ""]`.
+- To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `Float[Array, "..."]`.
+- You cannot have more than one use of multiple-axes, i.e. you can only use `...` or `*name` at most once in each array.
+- An example of broadcasting multiple dimensions:  
+    `def add(x: Float[Array, "*#foo"], y: Float[Array, "*#foo"]) -> Float[Array, "*#foo"]`.
+- A symbolic expression cannot be evaluated unless all of the axes sizes it refers to have already been processed. In practice this usually means that they should only be used in annotations for the return type, and only use axes declared in the arguments.
+
+## Dtype
+
+The dtype should be any one of (all imported from `jaxtyping`):
+
+- Any dtype at all: `Shaped`
+  - Boolean: `Bool`
+  - PRNG key: `Key`
+  - Any integer, unsigned integer, floating, or complex: `Num`
+    - Any floating or complex: `Inexact`
+      - Any floating point: `Float`
+        - Of particular precision: `BFloat16`, `Float16`, `Float32`, `Float64`
+      - Any complex: `Complex`
+        - Of particular precision: `Complex64`, `Complex128`
+    - Any integer or unsigned intger: `Integer`
+      - Any unsigned integer: `UInt`
+        - Of particular precision: `UInt8`, `UInt16`, `UInt32`, `UInt64`
+      - Any signed integer: `Int`
+        - Of particular precision: `Int8`, `Int16`, `Int32`, `Int64`
+
+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 Array, Float
+Float[Array, "some_shape"]
+```
+rather than
+```python
+from jaxtyping import Array, Float32
+Float32[Array, "some_shape"]
+```
+
+## Array
+
+The array should usually be a `jaxtyping.Array`, which is an alias for `jax.numpy.ndarray` (which is itself an alias for `jax.Array`).
+
+`jaxtyping.ArrayLike` is also available, which is an alias for `jax.typing.ArrayLike`. This is a union over JAX arrays and the builtin `bool`/`int`/`float`/`complex`.
+
+You can use non-JAX types as well. jaxtyping also supports NumPy, TensorFlow, and PyTorch, e.g.:
+```python
+Float[np.ndarray, "..."]
+Float[tf.Tensor, "..."]
+Float[torch.Tensor, "..."]
+```
+
+Shape-and-dtype specified jaxtyping arrays can also be used, e.g.
+```python
+Image = Float[Array, "channels height width"]
+BatchImage = Float[Image, "batch"]
+```
+in which case the additional shape is prepended, and the acceptable dtypes are the intersection of the two dtype specifiers used. (So that e.g. `BatchImage = Shaped[Image, "batch"]` would work just as well. But `Bool[Image, "batch"]` would throw an error, as there are no dtypes that are both bools and floats.) Thus the above is equivalent to
+```python
+BatchImage = Float[Array, "batch channels height width"]
+```
+
+Note that `jaxtyping.{Array, ArrayLike}` are only available if JAX has been installed.
+
+## Scalars, PRNGKeys
+
+For convenience, jaxtyping also includes `jaxtyping.Scalar`, `jaxtyping.ScalarLike`, and `jaxtyping.PRNGKeyArray`, defined as:
+```python
+Scalar = Shaped[Array, ""]
+ScalarLike = Shaped[ArrayLike, ""]
+
+# Depending on the value of `JAX_ENABLE_CUSTOM_PRNG`:
+PRNGKeyArray = Key[Array, ""]
+PRNGKeyArray = UInt32[Array, "2"]
+```
+
+Recalling that shape-and-dtype specified jaxtyping arrays can be nested, this means that e.g. you can annotate the output of `jax.random.split` with `Shaped[PRNGKeyArray, "2"]`, or e.g. an integer scalar with `Int[Scalar, ""]`.
+
+Note that `jaxtyping.{Scalar, ScalarLike, PRNGKeyArray}` are only available if JAX has been installed.

docs/api/pytree.md

@@ -0,0 +1,8 @@
+# PyTree annotations
+
+:::jaxtyping.PyTree
+    selection:
+        members:
+            false
+
+Note that `jaxtyping.PyTree` is only available if JAX has been installed.

docs/api/runtime-type-checking.md

@@ -0,0 +1,22 @@
+# Runtime type checking
+
+(See the [FAQ](../faq.md) for details on static type checking.)
+
+Runtime type checking **synergises beautifully with `jax.jit`!** All shape checks will be performed at trace-time only, and will not impact runtime performance.
+
+Runtime type-checking should be performed using a library like [typeguard](https://github.com/agronholm/typeguard) or [beartype](https://github.com/beartype/beartype).
+
+The types provided by `jaxtyping`, e.g. `Float[Array, "batch channels"]`, are all compatible with `isinstance` checks, e.g. `isinstance(x, Float[Array, "batch channels"])`. This means that jaxtyping should be compatible with all runtime type checkers out-of-the-box.
+
+Some additional context is needed to ensure consistency between multiple argments (i.e. that shapes match up between arrays). For this, you can use either `jaxtyping.jaxtyped` to add this capability to a single function, or `jaxtyping.install_import_hook` to add this capability to a whole codebase. If either are too much magic for you, you can safely use neither and have just single-argument type checking.
+
+::: jaxtyping.jaxtyped
+
+---
+
+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 hook.
+
+::: jaxtyping.install_import_hook

docs/faq.md

@@ -0,0 +1,48 @@
+# FAQ
+
+## Is jaxtyping compatible with static type checkers like `mypy`/`pyright`/`pytype`?
+
+There is partial support for these. An annotation of the form `dtype[array, shape]` should be treated as just `array` by a static type checker. Unfortunately full dtype/shape checking is beyond the scope of what static type checking is currently capable of.
+
+(Note that at time of writing, `pytype` has a bug in that `dtype[array, shape]` is sometimes treated as `Any` rather than `array`. `mypy` and `pyright` both work fine.)
+
+## How does jaxtyping interact with `jax.jit`?
+
+jaxtyping and `jax.jit` synergise beautifully.
+
+When calling JAX operations wrapped in a `jax.jit`, then the dtype/shape-checking will happen at trace time. (When JAX traces your function prior to compiling it.) The actual compiled code does not have any dtype/shape-checking, and will therefore still be just as fast as before!
+
+## `flake8` or Ruff are 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`, or Ruff, this can be resolved. Multi-dimensional arrays (e.g. `Float32[Array, "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. `Float32[Array, "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. `Float32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before.
+
+## Does jaxtyping use [PEP 646](https://www.python.org/dev/peps/pep-0646/) (variadic generics)?
+
+The intention of PEP 646 was to make it possible for static type checkers to perform shape checks of arrays. Unfortunately, this still isn't yet practical, so jaxtyping deliberately does not use this. (Yet?)
+
+The real problem is that Python's static typing ecosystem is a complicated collection of edge cases. Many of them block ML/scientific computing in particular. For example:
+
+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.
+
+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 were contemplated.)
+
+5. The syntax for static typing is a little verbose. You have to write things like `Array[Float32, Unpack[AnyShape], Literal[3], Height, Width]` instead of `Float32[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.

docs/index.md

@@ -0,0 +1,42 @@
+# Getting started
+
+jaxtyping is a library providing type annotations **and runtime type-checking** for:
+
+1. shape and dtype of [JAX](https://github.com/google/jax) arrays;
+2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
+
+ *(Now also supports PyTorch, NumPy, and TensorFlow!)*
+
+## Installation
+
+```bash
+pip install jaxtyping
+```
+
+Requires Python 3.8+.
+
+JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/etc.
+
+The annotations provided by jaxtyping are compatible with runtime type-checking packages, so it is common to also install one of these. 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).
+
+## Example
+
+```python
+from jaxtyping import Array, Float, PyTree
+
+# Accepts floating-point 2D arrays with matching dimensions
+def matrix_multiply(x: Float[Array, "dim1 dim2"],
+                    y: Float[Array, "dim2 dim3"]
+                  ) -> Float[Array, "dim1 dim3"]:
+    ...
+
+def accepts_pytree_of_ints(x: PyTree[int]):
+    ...
+
+def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
+    ...
+```
+
+## Next steps
+
+Have a read of the [Array annotations](./api/array.md) documentation on the left-hand bar!

docs/requirements.txt

@@ -0,0 +1,13 @@
+mkdocs==1.3.0            # Main documentation generator.
+mkdocs-material==7.3.6   # Theme
+pymdown-extensions==9.4  # Markdown extensions e.g. to handle LaTeX.
+mkdocstrings==0.17.0     # Autogenerate documentation from docstrings.
+mknotebooks==0.7.1       # Turn Jupyter Lab notebooks into webpages.
+pytkdocs_tweaks==0.0.5   # Tweaks mkdocstrings to improve various aspects
+mkdocs_include_exclude_files==0.0.1  # Tweak which files are included/excluded
+jinja2==3.0.3            # Older version. After 3.1.0 seems to be incompatible with current versions of mkdocstrings.
+pygments==2.14.0
+
+# Dependencies of jaxtyping itself.
+# Always use most up-to-date versions.
+jax[cpu]

jaxtyping/__init__.py

@@ -17,38 +17,150 @@
 # 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 importlib.metadata
+import typing
+
+
+try:
+    import jax
+except ImportError:
+    has_jax = False
+else:
+    has_jax = True
+    del jax
+
+# First import some things as normal
 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,
+    AbstractArray as AbstractArray,
+    AbstractDtype as AbstractDtype,
+    get_array_name_format as get_array_name_format,
+    set_array_name_format as set_array_name_format,
 )
-from .decorator import jaxtyped
-from .import_hook import install_import_hook
-from .pytree_type import PyTree
+from .decorator import jaxtyped as jaxtyped
+from .import_hook import install_import_hook as install_import_hook
 
 
-__version__ = "0.0.2"
+# Now import Array and ArrayLike
+if typing.TYPE_CHECKING:
+    # For imports, we need to explicitly `import X as X` in order for Pyright to see
+    # them as public. See discussion at https://github.com/microsoft/pyright/issues/2277
+    from jax import Array as Array
+    from jax.typing import ArrayLike as ArrayLike
+elif has_jax:
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+        class Array:
+            pass
+
+        Array.__module__ = "builtins"
+
+        class ArrayLike:
+            pass
+
+        ArrayLike.__module__ = "builtins"
+    else:
+        from jax import Array as Array
+
+        try:
+            from jax.typing import ArrayLike as ArrayLike
+        except (ModuleNotFoundError, ImportError):
+            pass
+
+
+# Import our dtypes
+if typing.TYPE_CHECKING:
+    # Introduce an indirection so that we can `import X as X` to make it clear that
+    # these are public.
+    from .indirection import (
+        BFloat16 as BFloat16,
+        Bool as Bool,
+        Complex as Complex,
+        Complex64 as Complex64,
+        Complex128 as Complex128,
+        Float as Float,
+        Float16 as Float16,
+        Float32 as Float32,
+        Float64 as Float64,
+        Inexact as Inexact,
+        Int as Int,
+        Int8 as Int8,
+        Int16 as Int16,
+        Int32 as Int32,
+        Int64 as Int64,
+        Integer as Integer,
+        Key as Key,
+        Num as Num,
+        Shaped as Shaped,
+        UInt as UInt,
+        UInt8 as UInt8,
+        UInt16 as UInt16,
+        UInt32 as UInt32,
+        UInt64 as UInt64,
+    )
+else:
+    from .array_types import (
+        BFloat16 as BFloat16,
+        Bool as Bool,
+        Complex as Complex,
+        Complex64 as Complex64,
+        Complex128 as Complex128,
+        Float as Float,
+        Float16 as Float16,
+        Float32 as Float32,
+        Float64 as Float64,
+        Inexact as Inexact,
+        Int as Int,
+        Int8 as Int8,
+        Int16 as Int16,
+        Int32 as Int32,
+        Int64 as Int64,
+        Integer as Integer,
+        Key as Key,
+        Num as Num,
+        Shaped as Shaped,
+        UInt as UInt,
+        UInt8 as UInt8,
+        UInt16 as UInt16,
+        UInt32 as UInt32,
+        UInt64 as UInt64,
+    )
+
+
+# Now import PyTree
+if typing.TYPE_CHECKING:
+    # Set up to deliberately confuse a static type checker.
+    import typing_extensions
+
+    PyTree: typing_extensions.TypeAlias = getattr(typing, "foo" + "bar")
+    # What's going on with this madness?
+    #
+    # At static-type-checking-time, we want `PyTree` to be a type for which both
+    # `PyTree` and `PyTree[Foo]` are equivalent to `Any`.
+    # (The intention is that `PyTree` be a runtime-only type; there's no real way to
+    # do more with static type checkers.)
+    #
+    # Unfortunately, this isn't possible: `Any` isn't subscriptable. And there's no
+    # equivalent way we can fake this using typing annotations. (In some sense the
+    # closest thing would be a `Protocol[T]` with no methods, but that's actually the
+    # opposite of what we want: that ends up allowing nothing at all.)
+    #
+    # The good news for us is that static type checkers have an internal escape hatch.
+    # If they can't figure out what a type is, then they just give up and allow
+    # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking
+    # annotation, which static type checkers aren't smart enough to resolve.
+elif has_jax:
+    from .pytree_type import PyTree as PyTree  # noqa: F401
+
+
+# Conveniences
+if typing.TYPE_CHECKING:
+    from jax.random import PRNGKeyArray as PRNGKeyArray
+
+    from .indirection import Scalar as Scalar, ScalarLike as ScalarLike
+elif has_jax:
+    from .array_types import PRNGKeyArray, Scalar, ScalarLike  # noqa: F401
+
+del has_jax
+
+
+__version__ = importlib.metadata.version("jaxtyping")

jaxtyping/array_types.py

@@ -17,15 +17,36 @@
 # 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 enum
 import functools as ft
-from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union
-from typing_extensions import Literal
+import re
+import sys
+import types
+import typing
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    NoReturn,
+    Optional,
+    Tuple,
+    Union,
+)
 
-import jax.numpy as jnp
+import numpy as np
 
 from .decorator import storage
 
 
+try:
+    import jax
+except ImportError:
+    has_jax = False
+else:
+    has_jax = True
+
+
 _array_name_format = "dtype_and_shape"
 
 
@@ -40,10 +61,17 @@ def set_array_name_format(value):
 
 _any_dtype = object()
 
+
 _anonymous_dim = object()
 _anonymous_variadic_dim = object()
 
 
+class _DimType(enum.Enum):
+    named = enum.auto()
+    fixed = enum.auto()
+    symbolic = enum.auto()
+
+
 class _NamedDim:
     def __init__(self, name, broadcastable):
         self.name = name
@@ -62,21 +90,28 @@ class _FixedDim:
         self.broadcastable = broadcastable
 
 
+class _SymbolicDim:
+    def __init__(self, expr, broadcastable):
+        self.expr = expr
+        self.broadcastable = broadcastable
+
+
 _AbstractDimOrVariadicDim = Union[
     Literal[_anonymous_dim],
     Literal[_anonymous_variadic_dim],
     _NamedDim,
     _NamedVariadicDim,
     _FixedDim,
+    _SymbolicDim,
 ]
-_AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim]
+_AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim, _SymbolicDim]
 
 
 def _check_dims(
     cls_dims: List[_AbstractDim],
     obj_shape: Tuple[int],
-    memo: Dict[str, Union[int, Tuple[int]]],
-):
+    single_memo: Dict[str, int],
+) -> bool:
     assert len(cls_dims) == len(obj_shape)
     for cls_dim, obj_size in zip(cls_dims, obj_shape):
         if cls_dim is _anonymous_dim:
@@ -86,12 +121,24 @@ def _check_dims(
         elif type(cls_dim) is _FixedDim:
             if cls_dim.size != obj_size:
                 return False
+        elif type(cls_dim) is _SymbolicDim:
+            try:
+                eval_size = eval(cls_dim.expr, single_memo)
+            except NameError as e:
+                raise NameError(
+                    f"Cannot process symbolic dimension '{cls_dim.expr}' as some "
+                    "dimension names have not been processed. In practice you should "
+                    "usually only use symbolic dimensions in annotations for return "
+                    "types, referring only to dimensions annotated for arguments."
+                ) from e
+            if eval_size != obj_size:
+                return False
         else:
             assert type(cls_dim) is _NamedDim
             try:
-                cls_size = memo[cls_dim.name]
+                cls_size = single_memo[cls_dim.name]
             except KeyError:
-                memo[cls_dim.name] = obj_size
+                single_memo[cls_dim.name] = obj_size
             else:
                 if cls_size != obj_size:
                     return False
@@ -100,36 +147,80 @@ def _check_dims(
 
 class _MetaAbstractArray(type):
     def __instancecheck__(cls, obj):
-        if not isinstance(obj, jnp.ndarray):
+        if not isinstance(obj, cls.array_type):
             return False
 
-        if cls.dtypes is not _any_dtype and obj.dtype not in cls.dtypes:
-            return False
+        if has_jax and jax.core.is_opaque_dtype(obj.dtype):
+            dtype = str(obj.dtype)
+        elif hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
+            # JAX, numpy
+            dtype = obj.dtype.type.__name__
+        elif hasattr(obj.dtype, "as_numpy_dtype"):
+            # TensorFlow
+            dtype = obj.dtype.as_numpy_dtype.__name__
+        else:
+            # PyTorch
+            repr_dtype = repr(obj.dtype).split(".")
+            if len(repr_dtype) == 2 and repr_dtype[0] == "torch":
+                dtype = repr_dtype[1]
+            else:
+                raise RuntimeError(
+                    "Unrecognised array/tensor type to extract dtype from"
+                )
 
-        if len(storage.memo_stack) == 0:
+        if cls.dtypes is not _any_dtype:
+            in_dtypes = False
+            for cls_dtype in cls.dtypes:
+                if type(cls_dtype) is str:
+                    in_dtypes = dtype == cls_dtype
+                elif type(cls_dtype) is re.Pattern:
+                    in_dtypes = bool(cls_dtype.match(dtype))
+                else:
+                    assert False
+                if in_dtypes:
+                    break
+            if not in_dtypes:
+                return False
+
+        no_temp_memo = hasattr(storage, "memo_stack") and len(storage.memo_stack) != 0
+
+        if no_temp_memo:
+            single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]
+            # Make a copy so we don't mutate the original memo during the shape check.
+            single_memo = single_memo.copy()
+            variadic_memo = variadic_memo.copy()
+            variadic_broadcast_memo = variadic_broadcast_memo.copy()
+        else:
             # `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
+            single_memo = {}
+            variadic_memo = {}
+            variadic_broadcast_memo = {}
 
-        if cls._check_shape(obj, memo):
+        if cls._check_shape(obj, single_memo, variadic_memo, variadic_broadcast_memo):
             # We update the memo every time we successfully pass a shape check
-            if not temp_memo:
-                storage.memo_stack[-1] = memo
+            if no_temp_memo:
+                storage.memo_stack[-1] = (
+                    single_memo,
+                    variadic_memo,
+                    variadic_broadcast_memo,
+                )
             return True
         else:
             return False
 
-    def _check_shape(cls, obj, memo):
+    def _check_shape(
+        cls,
+        obj,
+        single_memo: Dict[str, int],
+        variadic_memo: Dict[str, Tuple[int, ...]],
+        variadic_broadcast_memo: Dict[str, List[Tuple[int, ...]]],
+    ):
         if cls.index_variadic is None:
             if obj.ndim != len(cls.dims):
                 return False
-            return _check_dims(cls.dims, obj.shape, memo)
+            return _check_dims(cls.dims, obj.shape, single_memo)
         else:
             if obj.ndim < len(cls.dims) - 1:
                 return False
@@ -137,122 +228,369 @@ class _MetaAbstractArray(type):
             j = -(len(cls.dims) - i - 1)
             if j == 0:
                 j = None
-            if not _check_dims(cls.dims[:i], obj.shape[:i], memo):
+            if not _check_dims(cls.dims[:i], obj.shape[:i], single_memo):
                 return False
-            if j is not None and not _check_dims(cls.dims[j:], obj.shape[j:], memo):
+            if j is not None and not _check_dims(
+                cls.dims[j:], obj.shape[j:], single_memo
+            ):
                 return False
             variadic_dim = cls.dims[i]
-            if variadic_dim is not _anonymous_variadic_dim:
+            if variadic_dim is _anonymous_variadic_dim:
+                return True
+            else:
+                assert type(variadic_dim) is _NamedVariadicDim
                 variadic_name = variadic_dim.name
                 try:
-                    variadic_shape = memo[variadic_name]
+                    if variadic_dim.broadcastable:
+                        variadic_shapes = variadic_broadcast_memo[variadic_name]
+                    else:
+                        variadic_shape = variadic_memo[variadic_name]
                 except KeyError:
-                    memo[variadic_name] = obj.shape[i:j]
+                    if variadic_dim.broadcastable:
+                        variadic_broadcast_memo[variadic_name] = [obj.shape[i:j]]
+                    else:
+                        variadic_memo[variadic_name] = obj.shape[i:j]
+                    return True
                 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)
+                        new_shape = obj.shape[i:j]
+                        for existing_shape in variadic_shapes:
+                            try:
+                                np.broadcast_shapes(new_shape, existing_shape)
+                            except ValueError:
+                                return False
+                        variadic_shapes.append(new_shape)
+                        return True
                     else:
                         return variadic_shape == obj.shape[i:j]
-            return True
+            assert False
+
+
+@ft.lru_cache(maxsize=None)
+def _make_metaclass(base_metaclass):
+    class MetaAbstractArray(_MetaAbstractArray, base_metaclass):
+        pass
+
+    return MetaAbstractArray
+
+
+def _check_scalar(dtype, dtypes, dims):
+    if len(dims) != 0:
+        return dims == (_anonymous_variadic_dim,)
+    return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes)
 
 
 class AbstractArray(metaclass=_MetaAbstractArray):
-    dtypes: List[jnp.dtype]
-    dims: List[_AbstractDimOrVariadicDim]
+    """This is the base class of all shape-and-dtype-specified arrays, e.g. it's a base
+    class for `Float32[Array, "foo"]`.
+
+    This might be useful if you're trying to inspect type annotations yourself, e.g.
+    you can check `issubclass(annotation, jaxtyping.AbstractArray)`.
+    """
+
+    array_type: Any
+    dtypes: List[str]
+    dims: Tuple[_AbstractDimOrVariadicDim, ...]
     index_variadic: Optional[int]
+    dim_str: str
+
+
+_not_made = object()
+
+
+_union_types = [typing.Union]
+if sys.version_info >= (3, 10):
+    _union_types.append(types.UnionType)
+
+
+@ft.lru_cache(maxsize=None)
+def _make_array(array_type, dim_str, dtypes, name):
+    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 and "(" not in elem:
+            # Common mistake.
+            # Disable in the case that there's brackets to allow for function calls,
+            # e.g. `min(foo,bar)`, in symbolic dimensions.
+            raise ValueError("Dimensions should be separated with spaces, not commas")
+        if elem.endswith("#"):
+            raise ValueError(
+                "As of jaxtyping v0.1.0, broadcastable dimensions are now denoted "
+                "with a # at the start, rather than at the end"
+            )
+
+        if "..." in elem:
+            if elem != "...":
+                raise ValueError(
+                    "Anonymous multiple dimension '...' must be used on its own; "
+                    f"got {elem}"
+                )
+            broadcastable = False
+            variadic = True
+            anonymous = True
+            dim_type = _DimType.named
+        else:
+            broadcastable = False
+            variadic = False
+            anonymous = False
+            while True:
+                if len(elem) == 0:
+                    # This branch needed as just `_` is valid
+                    break
+                first_char = elem[0]
+                if first_char == "#":
+                    if broadcastable:
+                        raise ValueError(
+                            "Do not use # twice to denote broadcastability, e.g. "
+                            "`##foo` is not allowed"
+                        )
+                    broadcastable = True
+                    elem = elem[1:]
+                elif first_char == "*":
+                    if variadic:
+                        raise ValueError(
+                            "Do not use * twice to denote accepting multiple "
+                            "dimensions, e.g. `**foo` is not allowed"
+                        )
+                    variadic = True
+                    elem = elem[1:]
+                elif first_char == "_":
+                    if anonymous:
+                        raise ValueError(
+                            "Do not use _ twice to denote anonymity, e.g. `__foo` "
+                            "is not allowed"
+                        )
+                    anonymous = True
+                    elem = elem[1:]
+                # Allow e.g. `foo=4` as an alternate syntax for just `4`, so that one
+                # can write e.g. `Float[Array, "rows=3 cols=4"]`
+                elif elem.count("=") == 1:
+                    _, elem = elem.split("=")
+                else:
+                    break
+            if len(elem) == 0 or elem.isidentifier():
+                dim_type = _DimType.named
+            else:
+                try:
+                    elem = int(elem)
+                except ValueError:
+                    dim_type = _DimType.symbolic
+                else:
+                    dim_type = _DimType.fixed
+
+        if variadic:
+            if index_variadic is not None:
+                raise ValueError(
+                    "Cannot use multiple-dimension specifiers (`*name` or `...`) "
+                    "more than once"
+                )
+            index_variadic = index
+
+        if dim_type is _DimType.fixed:
+            if variadic:
+                raise ValueError(
+                    "Cannot have a fixed axis bind to multiple dimensions, e.g. "
+                    "`*4` is not allowed"
+                )
+            if anonymous:
+                raise ValueError(
+                    "Cannot have a fixed axis be anonymous, e.g. `_4` is not " "allowed"
+                )
+            elem = _FixedDim(elem, broadcastable)
+        elif dim_type is _DimType.named:
+            if anonymous:
+                if broadcastable:
+                    raise ValueError(
+                        "Cannot have a dimension be both anonymous and "
+                        "broadcastable, e.g. `#_` is not allowed"
+                    )
+                if variadic:
+                    elem = _anonymous_variadic_dim
+                else:
+                    elem = _anonymous_dim
+            else:
+                if variadic:
+                    elem = _NamedVariadicDim(elem, broadcastable)
+                else:
+                    elem = _NamedDim(elem, broadcastable)
+        else:
+            assert dim_type is _DimType.symbolic
+            if anonymous:
+                raise ValueError(
+                    "Cannot have a symbolic dimension be anonymous, e.g. "
+                    "`_foo+bar` is not allowed"
+                )
+            if variadic:
+                raise ValueError(
+                    "Cannot have symbolic multiple-dimensions, e.g. "
+                    "`*foo+bar` is not allowed"
+                )
+            elem = compile(elem, "<string>", "eval")
+            elem = _SymbolicDim(elem, broadcastable)
+        dims.append(elem)
+    dims = tuple(dims)
+
+    # Allow Python built-in numeric types.
+    # TODO: do something more generic than this? Should we _make all types
+    # that have `shape` and `dtype` attributes or something?
+    if array_type is bool:
+        if _check_scalar("bool", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is int:
+        if _check_scalar("int", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is float:
+        if _check_scalar("float", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is complex:
+        if _check_scalar("complex", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    if issubclass(array_type, AbstractArray):
+        if dtypes is _any_dtype:
+            dtypes = array_type.dtypes
+        elif array_type.dtypes is not _any_dtype:
+            dtypes = tuple(x for x in dtypes if x in array_type.dtypes)
+            if len(dtypes) == 0:
+                raise ValueError(
+                    "A jaxtyping annotation cannot be extended with no overlapping "
+                    "dtypes. For example, `Bool[Float[Array, 'dim1'], 'dim2']` is an "
+                    "error. You probably want to make the outer wrapper be `Shaped`."
+                )
+        if array_type.index_variadic is not None:
+            if index_variadic is None:
+                index_variadic = array_type.index_variadic + len(dims)
+            else:
+                raise ValueError(
+                    "Cannot use multiple-dimension specifiers (`*name` or `...`) "
+                    "in both the original array and the extended array"
+                )
+        dims = dims + array_type.dims
+        dim_str = dim_str + " " + array_type.dim_str
+        array_type = array_type.array_type
+    try:
+        type_str = array_type.__name__
+    except AttributeError:
+        type_str = repr(array_type)
+    if _array_name_format == "dtype_and_shape":
+        name = f"{name}[{type_str}, '{dim_str}']"
+    elif _array_name_format == "array":
+        name = type_str
+    else:
+        raise ValueError(f"array_name_format {_array_name_format} not recognised")
+    metaclass = _make_metaclass(type(array_type))
+    out = metaclass(
+        name,
+        (array_type, AbstractArray),
+        dict(
+            array_type=array_type,
+            dtypes=dtypes,
+            dims=dims,
+            index_variadic=index_variadic,
+            dim_str=dim_str,
+        ),
+    )
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        out.__module__ = "builtins"
+    else:
+        out.__module__ = "jaxtyping"
+    return out
 
 
 class _MetaAbstractDtype(type):
     def __instancecheck__(cls, obj: Any) -> NoReturn:
         raise RuntimeError(
-            f"Do not use `isinstance(x, jaxtyping.{cls.__name__}`. If you want to "
+            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__}["..."]`.'
+            f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
         )
 
-    @ft.lru_cache(maxsize=None)
-    def __getitem__(cls, dim_str: str) -> _MetaAbstractArray:
-        if not isinstance(dim_str, str):
+    def __getitem__(cls, item: Tuple[Any, str]):
+        if not isinstance(item, tuple) or len(item) != 2:
             raise ValueError(
-                "Shape specification must be a string. Axes should be separated with spaces."
+                "As of jaxtyping v0.2.0, type annotations must now include an explicit "
+                "array type. For example `jaxtyping.Float32[jax.Array, 'foo bar']`."
             )
-        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)
+        array_type, dim_str = item
+        del item
+        if typing.get_origin(array_type) in _union_types:
+            out = [
+                _make_array(x, dim_str, cls.dtypes, cls.__name__)
+                for x in typing.get_args(array_type)
+            ]
+            out = tuple(x for x in out if x is not _not_made)
+            if len(out) == 0:
+                raise ValueError("Invalid jaxtyping type annotation.")
+            elif len(out) == 1:
+                (out,) = out
             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"
+                out = Union[out]
         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),
-        )
+            out = _make_array(array_type, dim_str, cls.dtypes, cls.__name__)
+            if out is _not_made:
+                raise ValueError("Invalid jaxtyping type annotation.")
+        return out
 
 
 class AbstractDtype(metaclass=_MetaAbstractDtype):
-    dtypes: Union[str, List[str], Literal[_any_dtype]]
+    """This is the base class of all dtypes. This can be used to create your own custom
+    collection of dtypes (analogous to `Float`, `Inexact` etc.)
+
+    You must specify the class attribute `dtypes`. This can either be a string, a
+    regex (as returned by `re.compile(...)`), or a tuple/list of strings/regexes.
+
+    At runtime, the array or tensor's dtype is converted to a string and compared
+    against the string (an exact match is required) or regex. (String matching is
+    performed, rather than just e.g. `array.dtype == dtype`, to provide cross-library
+    compatibility between JAX/PyTorch/etc.)
+
+    !!! Example
+
+        ```python
+        class UInt8or16(AbstractDtype):
+            dtypes = ["uint8", "uint16"]
+
+        UInt8or16[Array, "shape"]
+        ```
+        which is essentially equivalent to
+        ```python
+        Union[UInt8[Array, "shape"], UInt16[Array, "shape"]]
+        ```
+    """
+
+    dtypes: Union[Literal[_any_dtype], List[Union[str, re.Pattern]]]
 
     def __init__(self, *args, **kwargs):
         raise RuntimeError(
             "AbstractDtype cannot be instantiated. Perhaps you wrote e.g. "
-            '`f32("shape")` when you mean `f32["shape"]`?'
+            '`Float32("shape")` when you mean `Float32[jnp.ndarray, "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]
+        dtypes: Union[Literal[_any_dtype], str, List[str]] = cls.dtypes
+        if isinstance(dtypes, (str, re.Pattern)):
+            dtypes = (dtypes,)
+        elif dtypes is not _any_dtype:
+            dtypes = tuple(dtypes)
         cls.dtypes = dtypes
 
 
 _bool = "bool"
+_bool_ = "bool_"
 _uint8 = "uint8"
 _uint16 = "uint16"
 _uint32 = "uint32"
@@ -275,25 +613,29 @@ def _make_dtype(_dtypes, name):
 
     _Cls.__name__ = name
     _Cls.__qualname__ = name
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        _Cls.__module__ = "builtins"
+    else:
+        _Cls.__module__ = "jaxtyping"
     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")
+UInt8 = _make_dtype(_uint8, "UInt8")
+UInt16 = _make_dtype(_uint16, "UInt16")
+UInt32 = _make_dtype(_uint32, "UInt32")
+UInt64 = _make_dtype(_uint64, "UInt64")
+Int8 = _make_dtype(_int8, "Int8")
+Int16 = _make_dtype(_int16, "Int16")
+Int32 = _make_dtype(_int32, "Int32")
+Int64 = _make_dtype(_int64, "Int64")
+BFloat16 = _make_dtype(_bfloat16, "BFloat16")
+Float16 = _make_dtype(_float16, "Float16")
+Float32 = _make_dtype(_float32, "Float32")
+Float64 = _make_dtype(_float64, "Float64")
+Complex64 = _make_dtype(_complex64, "Complex64")
+Complex128 = _make_dtype(_complex128, "Complex128")
 
+bools = [_bool, _bool_]
 uints = [_uint8, _uint16, _uint32, _uint64]
 ints = [_int8, _int16, _int32, _int64]
 floats = [_bfloat16, _float16, _float32, _float64]
@@ -301,15 +643,25 @@ 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")
+Bool = _make_dtype(bools, "Bool")
+UInt = _make_dtype(uints, "UInt")
+Int = _make_dtype(ints, "Int")
+Integer = _make_dtype(uints + ints, "Integer")
+Float = _make_dtype(floats, "Float")
+Complex = _make_dtype(complexes, "Complex")
+Inexact = _make_dtype(floats + complexes, "Inexact")
+Num = _make_dtype(uints + ints + floats + complexes, "Num")
+
+Shaped = _make_dtype(_any_dtype, "Shaped")
+
+if has_jax:
+    if jax.config.jax_enable_custom_prng:
+        _key_regex = re.compile(r"^key<\w+>$")
+        Key = _make_dtype(_key_regex, "Key")
+        PRNGKeyArray = Key[jax.Array, ""]
+    else:
+        Key = UInt32
+        PRNGKeyArray = Key[jax.Array, "2"]
+    Scalar = Shaped[jax.Array, ""]
+    ScalarLike = Shaped[jax.typing.ArrayLike, ""]

jaxtyping/decorator.py

@@ -17,22 +17,127 @@
 # 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 dataclasses
 import functools as ft
+import inspect
 import threading
+import types
+import weakref
 
 
 storage = threading.local()
-storage.memo_stack = []
+
+
+_jaxtyped_fns = weakref.WeakSet()
 
 
 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()
+    """Used in conjunction with a runtime type checker. Decorate a function with this to
+    have shapes checked for consistency across multiple arguments.
 
-    return wrapper
+    Note that `@jaxtyped` is applied above the type checker.
+
+    !!! Example
+
+        ```python
+        # Import both the annotation and the `jaxtyped` decorator from `jaxtyping`
+        from jaxtyping import Array, Float32, 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: Float32[Array, "b c1"],
+                                y: Float32[Array, "b c2"]
+                              ) -> Float32[Array, "b c1 c2"]:
+            return x[:, :, None] * y[:, None, :]
+        ```
+
+    **Notes 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 that analyses
+    arguments, and instead 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, Float32[Array, "foo"])`.
+    """
+    if type(fn) is types.FunctionType and fn in _jaxtyped_fns:
+        return fn
+    elif inspect.isclass(fn):  # allow decorators on class definitions
+        if dataclasses.is_dataclass(fn):
+            init = jaxtyped(fn.__init__)
+            fn.__init__ = init
+            return fn
+        else:
+            raise ValueError(
+                "jaxtyped may only be added as a class decorator to dataclasses"
+            )
+    # It'd be lovely if we could handle arbitrary descriptors, and not just the builtin
+    # ones. Unfortunately that means returning a class instance with a __get__ method,
+    # and that turns out to break loads of other things. See beartype issue #211 and
+    # jaxtyping issue #71.
+    elif isinstance(fn, classmethod):
+        return classmethod(jaxtyped(fn.__func__))
+    elif isinstance(fn, staticmethod):
+        return staticmethod(jaxtyped(fn.__func__))
+    elif isinstance(fn, property):
+        if fn.fget is None:
+            fget = None
+        else:
+            fget = jaxtyped(fn.fget)
+        if fn.fset is None:
+            fset = None
+        else:
+            fset = jaxtyped(fn.fset)
+        if fn.fdel is None:
+            fdel = None
+        else:
+            fdel = jaxtyped(fn.fdel)
+        return property(fget=fget, fset=fset, fdel=fdel)
+    else:
+
+        @ft.wraps(fn)
+        def wrapped_fn(*args, **kwargs):
+            try:
+                memo_stack = storage.memo_stack
+            except AttributeError:
+                memo_stack = storage.memo_stack = []
+            memo_stack.append(({}, {}, {}))
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                memo_stack.pop()
+
+        _jaxtyped_fns.add(wrapped_fn)
+        return wrapped_fn
+
+
+def _jaxtyped_typechecker(typechecker):
+    # typechecker is expected to probably be either `typeguard.typechecked`, or
+    # `beartype.beartype`, or `None`.
+
+    if typechecker is None:
+        typechecker = lambda x: x
+
+    def _wrapper(kls):
+        assert inspect.isclass(kls)
+        if dataclasses.is_dataclass(kls):
+            init = jaxtyped(typechecker(kls.__init__))
+            kls.__init__ = init
+        return kls
+
+    return _wrapper

jaxtyping/import_hook.py

@@ -31,30 +31,32 @@
 #
 # 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:
+# 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 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.
+# 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 functools as ft
 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 typing import List, Optional, Sequence, Union
 from unittest.mock import patch
 
 
@@ -63,8 +65,31 @@ 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")
+def _optimized_cache_from_source(typechecker_hash, /, path, debug_override=None):
+    # Version 2: change the position of the `@jaxtyped` decorator, so need a
+    #     different name to avoid hitting old __pycache__.
+    # Version 3: now also annotating classes.
+    # Version 4: I'm honestly not sure, but bumping this fixed some kind of odd error.
+    #     Maybe I changed something with hte classes part way through version 3?
+    # Version 5: Added support for string-based `typechecker` argument.
+    # Version 6: optimization tag now depends on `typechecker` argument, so that
+    #    changing the typechecker will hit a different cache.
+    return cache_from_source(
+        path, debug_override, optimization=f"jaxtyping6{typechecker_hash}"
+    )
+
+
+def _dot_lookup(*elements):
+    out = ast.Name(id=elements[0], ctx=ast.Load())
+    for element in elements[1:]:
+        out = ast.Attribute(out, element, ctx=ast.Load())
+    return out
+
+
+def _str_lookup(string):
+    module = ast.parse(string)
+    (expr,) = module.body
+    return expr.value
 
 
 class _JaxtypingTransformer(ast.NodeVisitor):
@@ -83,7 +108,7 @@ class _JaxtypingTransformer(ast.NodeVisitor):
             else:
                 node.body.insert(i, ast.Import(names=[ast.alias("jaxtyping", None)]))
                 if self._typechecker is not None:
-                    typechecker_module, _ = self._typechecker
+                    typechecker_module, _ = self._typechecker.split(".", 1)
                     node.body.insert(
                         i, ast.Import(names=[ast.alias(typechecker_module, None)])
                     )
@@ -94,31 +119,41 @@ class _JaxtypingTransformer(ast.NodeVisitor):
         self._parents.pop()
         return node
 
+    def visit_ClassDef(self, node: ast.ClassDef):
+        func = _dot_lookup("jaxtyping", "decorator", "_jaxtyped_typechecker")
+        if self._typechecker is None:
+            args = [ast.Constant(None)]
+        else:
+            args = [_str_lookup(self._typechecker)]
+        node.decorator_list.insert(0, ast.Call(func, args, keywords=[]))
+        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()
-                ),
-            )
+            # Place at the end of the decorator list, as otherwise we wrap e.g.
+            # `jax.custom_{jvp,vjp}` and lose the ability to `defjvp` etc.
+            #
+            # Note that the counter-argument here is that we'd like to place this
+            # at the start of the decorator list, in case a typechecking annotation
+            # has been manually applied, and we'd need to be above that. In this
+            # case we're just going to have to need to ask the user to remove their
+            # typechecking annotation (and let this decorator do it instead).
+            # It's more important we be compatible with normal JAX code.
+            #
+            #
+            # FWIW, typeguard also wants to be at the end of the decorator list, as it
+            # works by recompiling the wrapped function.
+            node.decorator_list.append(_dot_lookup("jaxtyping", "jaxtyped"))
             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(),
-                    )
-                )
-
+                node.decorator_list.append(_str_lookup(self._typechecker))
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
@@ -129,6 +164,7 @@ class _JaxtypingLoader(SourceFileLoader):
     def __init__(self, *args, typechecker, **kwargs):
         super().__init__(*args, **kwargs)
         self._typechecker = typechecker
+        self._typechecker_hash = str(abs(hash(self._typechecker)))
 
     def source_to_code(self, data, path, *, _optimize=-1):
         source = decode_source(data)
@@ -148,10 +184,11 @@ class _JaxtypingLoader(SourceFileLoader):
         )
 
     def exec_module(self, module):
-        # Use a custom optimization marker – the import lock should make this monkey patch safe
+        # 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,
+            ft.partial(_optimized_cache_from_source, self._typechecker_hash),
         ):
             return super().exec_module(module)
 
@@ -213,50 +250,127 @@ class ImportHookManager:
 
 # 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.
+def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optional[str]):
+    """Automatically apply `@jaxtyped`, and optionally a type checker, as decorators.
 
-    It will only be applied to modules loaded **after** this hook has been installed.
+    !!! Tip "Usage"
+
+        ```python
+        from jaxtyping import install_import_hook
+        # Plus any one of the following:
+
+        # decorate @jaxtyped and @typeguard.typechecked
+        with install_import_hook("foo", "typeguard.typechecked"):
+            import foo          # Any module imported inside this `with` block, whose
+            import foo.bar      # name begins with the specified string, will
+            import foo.bar.qux  # automatically have both `@jaxtyped` and the specified
+                                # typechecker applied to all of their functions.
+
+        # decorate @jaxtyped and @beartype.beartype
+        with install_import_hook("foo", "beartype.beartype"):
+            ...
+
+        # decorate only @jaxtyped (if you want that for some reason)
+        with install_import_hook("foo", None):
+            ...
+        ```
+
+        If you don't like using the `with` block, the hook can be used without that:
+        ```python
+        hook = install_import_hook(...):
+        import ...
+        hook.uninstall()
+        ```
+
+        The import hook can be applied to multiple packages via
+        ```python
+        install_import_hook(["foo", "bar.baz"], ...)
+        ```
+
+    The import hook will automatically decorate all functions, and the `__init__` method
+    of dataclasses.
+
+    If the function already has any decorators on it, then both the `@jaxtyped` and the
+    typechecker decorators will get added at the bottom of the decorator list, e.g.
+    ```python
+    @some_other_decorator
+    @jaxtyped
+    @beartype.beartype
+    def foo(...): ...
+    ```
 
     **Arguments:**:
 
-    - `packages`: the names of the modules in which to automatically apply `@jaxtyped`
+    - `modules`: 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.
+        string. 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.
 
     **Returns:**
 
     A context manager that uninstalls the hook on exit, or when you call `.uninstall()`.
 
-    **Example:**
+    ??? Example "Example: end-user script"
 
-    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
+        ```python
+        ### entry_point.py
+        from jaxtyping import install_import_hook
+        with install_import_hook("main", "typeguard.typechecked"):
+            import main
 
-    # main.py
-    from jaxtyped import f32
+        ### main.py
+        from jaxtyping import Array, Float32
+
+        def f(x: Float32[Array, "batch channels"]):
+            ...
+        ```
+
+    ??? Example "Example: writing a library"
+
+        ```python
+        ### __init__.py
+        from jaxtyping import install_import_hook
+        with 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.
+        ```
+
+    ??? info "Pytest hook"
+
+        The import hook can be installed at test-time only, as a pytest hook. From the
+        command line the syntax is:
+        ```
+        pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
+        ```
+        or in `pyproject.toml`:
+        ```toml
+        [tool.pytest.ini_options]
+        addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
+        ```
+        or in `pytest.ini`:
+        ```ini
+        [pytest]
+        addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
+        ```
+        This example 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`.
+
+        (This is the author's preferred approach to performing runtime type-checking
+        with jaxtyping!)
+    """  # noqa: E501
 
-    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]
 
+    # Support old less-flexible API.
+    if isinstance(typechecker, tuple):
+        typechecker = ".".join(typechecker)
+
     for i, finder in enumerate(sys.meta_path):
         if (
             isclass(finder)

jaxtyping/indirection.py

@@ -0,0 +1,32 @@
+# Note that `from typing_extensions import Annotated; Bool = Annotated`
+# does not work with static type checkers. `Annotated` is a typeform rather
+# than a type, meaning it cannot be assigned.
+from typing_extensions import (
+    Annotated as BFloat16,  # noqa: F401
+    Annotated as Bool,  # noqa: F401
+    Annotated as Complex,  # noqa: F401
+    Annotated as Complex64,  # noqa: F401
+    Annotated as Complex128,  # noqa: F401
+    Annotated as Float,  # noqa: F401
+    Annotated as Float16,  # noqa: F401
+    Annotated as Float32,  # noqa: F401
+    Annotated as Float64,  # noqa: F401
+    Annotated as Inexact,  # noqa: F401
+    Annotated as Int,  # noqa: F401
+    Annotated as Int8,  # noqa: F401
+    Annotated as Int16,  # noqa: F401
+    Annotated as Int32,  # noqa: F401
+    Annotated as Int64,  # noqa: F401
+    Annotated as Integer,  # noqa: F401
+    Annotated as Key,  # noqa: F401
+    Annotated as Num,  # noqa: F401
+    Annotated as Shaped,  # noqa: F401
+    Annotated as UInt,  # noqa: F401
+    Annotated as UInt8,  # noqa: F401
+    Annotated as UInt16,  # noqa: F401
+    Annotated as UInt32,  # noqa: F401
+    Annotated as UInt64,  # noqa: F401
+)
+
+from jax import Array as Scalar  # noqa: F401
+from jax.typing import ArrayLike as ScalarLike  # noqa: F401

jaxtyping/py.typed

@@ -0,0 +1 @@
+

jaxtyping/pytest_plugin.py

@@ -52,4 +52,4 @@ def pytest_configure(config):
         )
         raise RuntimeError(message.format(", ".join(already_imported_packages)))
 
-    install_import_hook(packages, typechecker.rsplit(".", 1))
+    install_import_hook(packages, typechecker)

jaxtyping/pytree_type.py

@@ -18,9 +18,10 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import functools as ft
+import typing
 from typing import Generic, TypeVar
 
-import jax
+import jax.tree_util as jtu
 import typeguard
 
 
@@ -33,6 +34,7 @@ class _FakePyTree(Generic[_T]):
 
 _FakePyTree.__name__ = "PyTree"
 _FakePyTree.__qualname__ = "PyTree"
+_FakePyTree.__module__ = "builtins"
 # Can't do type("PyTree", (Generic[_T],), {}) because dynamic subclassing of typeforms
 # isn't allowed.
 # Can't do types.new_class("PyTree", (Generic[_T],), {}) because that has __module__
@@ -49,7 +51,20 @@ class _MetaPyTree(type):
     @ft.lru_cache(maxsize=None)
     def __getitem__(cls, item):
         name = str(_FakePyTree[item])
-        return _MetaSubscriptPyTree(name, (), {"leaftype": item})
+        out = _MetaSubscriptPyTree(name, (), {"leaftype": item})
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            out.__module__ = "builtins"
+        else:
+            out.__module__ = "jaxtyping"
+        return out
+
+
+try:
+    # new typeguard
+    _TypeCheckError = (TypeError, typeguard.TypeCheckError)
+except AttributeError:
+    # old typeguard
+    _TypeCheckError = TypeError
 
 
 class _MetaSubscriptPyTree(type):
@@ -70,16 +85,29 @@ class _MetaSubscriptPyTree(type):
         def is_leaftype(x):
             try:
                 accepts_leaftype(x)
-            except TypeError:
+            except _TypeCheckError:
                 return False
             else:
                 return True
 
-        leaves = jax.tree_leaves(obj, is_leaf=is_leaftype)
+        leaves = jtu.tree_leaves(obj, is_leaf=is_leaftype)
         return all(map(is_leaftype, leaves))
 
 
-PyTree = _MetaPyTree("PyTree", (), {})
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
-# instancecheck for PyTree[foo], but we subclassing
+# instancecheck for PyTree[foo], but subclassing
 # `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.
+PyTree = _MetaPyTree("PyTree", (), {})
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+    PyTree.__module__ = "builtins"
+else:
+    PyTree.__module__ = "jaxtyping"
+PyTree.__doc__ = """Represents a PyTree.
+
+Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or
+`PyTree[Union[str, Float32[Array, "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))
+"""  # noqa: E501

mkdocs.yml

@@ -0,0 +1,100 @@
+theme:
+    name: material
+    features:
+        - navigation.sections  # Sections are included in the navigation on the left.
+        - toc.integrate  # Table of contents is integrated on the left; does not appear separately on the right.
+        - header.autohide  # header disappears as you scroll
+    palette:
+        # Light mode / dark mode
+        # We deliberately don't automatically use `media` to check a user's preferences. We default to light mode as
+        # (a) it looks more professional, and (b) is more obvious about the fact that it offers a (dark mode) toggle.
+        - scheme: default
+          primary: white
+          accent: amber
+          toggle:
+             icon: material/weather-night
+             name: Switch to dark mode
+        - scheme: slate
+          primary: black
+          accent: amber
+          toggle:
+             icon: material/weather-sunny
+             name: Switch to light mode
+    icon:
+        repo: fontawesome/brands/github  # GitHub logo in top right
+        logo: "material/check-network-outline"  # jaxtyping logo in top left
+    favicon: "_static/favicon.png"
+    custom_dir: "docs/_overrides"  # Overriding part of the HTML
+
+    # These additions are my own custom ones, having overridden a partial.
+    twitter_name: "@PatrickKidger"
+    twitter_url: "https://twitter.com/PatrickKidger"
+
+site_name: jaxtyping
+site_description: The documentation for the jaxtyping software library.
+site_author: Patrick Kidger
+site_url: https://docs.kidger.site/jaxtyping
+
+repo_url: https://github.com/google/jaxtyping
+repo_name: google/jaxtyping
+edit_uri: ""  # No edit button, as some of our pages are in /docs and some in /examples via symlink, so it's impossible for them all to be accurate
+
+strict: true  # Don't allow warnings during the build process
+
+extra_javascript: 
+    # The below three make MathJax work, see https://squidfunk.github.io/mkdocs-material/reference/mathjax/
+    - _static/mathjax.js
+    - https://polyfill.io/v3/polyfill.min.js?features=es6
+    - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js
+
+extra_css:
+    - _static/custom_css.css
+
+markdown_extensions:
+    - pymdownx.arithmatex:  # Render LaTeX via MathJax
+        generic: true
+    - pymdownx.superfences  # Seems to enable syntax highlighting when used with the Material theme.
+    - pymdownx.details  # Allowing hidden expandable regions denoted by ???
+    - pymdownx.snippets:  # Include one Markdown file into another
+        base_path: docs
+    - admonition
+    - toc:
+        permalink: "¤"  # Adds a clickable permalink to each section heading
+        toc_depth: 4
+
+plugins:
+    - search  # default search plugin; needs manually re-enabling when using any other plugins
+    - autorefs  # Cross-links to headings
+    - include_exclude_files:
+        include:
+            - ".htaccess"
+        exclude:
+            - "_overrides"
+    - mknotebooks  # Jupyter notebooks
+    - mkdocstrings:
+        handlers:
+            python:
+                setup_commands:
+                    - import pytkdocs_tweaks, typing
+                    - pytkdocs_tweaks.main()
+                    - typing.GENERATING_DOCUMENTATION = False
+
+                selection:
+                    inherited_members: true  # Allow looking up inherited methods
+                rendering:
+                    show_root_heading: true  # actually display anything at all...
+                    show_root_full_path: true  # display "diffrax.asdf" not just "asdf"
+                    show_if_no_docstring: true
+                    show_signature_annotations: true
+                    show_source: false  # don't include source code
+                    members_order: source  # order methods according to their order of definition in the source code, not alphabetical order
+                    heading_level: 4
+
+nav:
+    - 'index.md'
+    - API:
+        - 'api/array.md'
+        - 'api/pytree.md'
+        - 'api/runtime-type-checking.md'
+        - 'api/advanced-features.md'
+    - 'faq.md'

pyproject.toml

@@ -0,0 +1,45 @@
+[project]
+name = "jaxtyping"
+version = "0.2.17"
+description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
+readme = "README.md"
+requires-python ="~=3.8"
+license = {file = "LICENSE"}
+authors = [
+  {name = "Patrick Kidger", email = "contact@kidger.site"},
+]
+keywords = ["jax", "neural-networks", "deep-learning", "equinox", "typing"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Financial and Insurance Industry",
+    "Intended Audience :: Information Technology",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Natural Language :: English",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Information Analysis",
+    "Topic :: Scientific/Engineering :: Mathematics",
+]
+urls = {repository = "https://github.com/google/jaxtyping" }
+dependencies = ["numpy>=1.20.0", "typeguard>=2.13.3", "typing_extensions>=3.7.4.1"]
+entry_points = {pytest11 = ["jaxtyping = jaxtyping.pytest_plugin"]}
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build]
+include = ["jaxtyping/*"]
+
+[tool.ruff]
+select = ["E", "F", "I001"]
+ignore = ["E721", "E731", "F722"]
+ignore-init-module-imports = true
+
+[tool.ruff.isort]
+combine-as-imports = true
+lines-after-imports = 2
+extra-standard-library = ["typing_extensions"]
+order-by-type = false

setup.py

@@ -1,88 +0,0 @@
-# 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", "typing_extensions>=4.2.0"]
-
-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],
-)

test/conftest.py

@@ -19,13 +19,24 @@
 
 import random
 
-import beartype
 import jax.random as jr
 import pytest
 import typeguard
 
 
-@pytest.fixture(params=[typeguard.typechecked, beartype.beartype])
+try:
+    import beartype
+except ImportError:
+
+    def skip(*args, **kwargs):
+        pytest.skip("Beartype not installed")
+
+    typecheck_params = [typeguard.typechecked, skip]
+else:
+    typecheck_params = [typeguard.typechecked, beartype.beartype]
+
+
+@pytest.fixture(params=typecheck_params)
 def typecheck(request):
     return request.param
 
@@ -37,6 +48,3 @@ def getkey():
         return jr.PRNGKey(random.randint(0, 2**31 - 1))
 
     return _getkey
-
-
-ParamException = (TypeError, beartype.roar.BeartypeCallHintParamViolation)

test/helpers.py

@@ -17,12 +17,32 @@
 # 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
+import typeguard
 
 
-ParamError = (TypeError, beartype.roar.BeartypeCallHintParamViolation)
-ReturnError = (TypeError, beartype.roar.BeartypeCallHintReturnViolation)
+ParamError = []
+ReturnError = []
+ParamError.append(TypeError)  # old typeguard
+ReturnError.append(TypeError)  # old typeguard
+
+try:
+    # new typeguard
+    ParamError.append(typeguard.TypeCheckError)
+    ReturnError.append(typeguard.TypeCheckError)
+except AttributeError:
+    pass
+
+try:
+    import beartype
+except ImportError:
+    pass
+else:
+    ParamError.append(beartype.roar.BeartypeCallHintParamViolation)
+    ReturnError.append(beartype.roar.BeartypeCallHintReturnViolation)
+
+ParamError = tuple(ParamError)
+ReturnError = tuple(ReturnError)
 
 
 @eqx.filter_jit

test/import_hook_tester_beartype.py

@@ -17,17 +17,46 @@
 # 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 dataclasses
+
+import equinox as eqx
 import jax.numpy as jnp
 import pytest
-from helpers import ParamError
 
-from jaxtyping import f32
+from jaxtyping import Float32
+
+from .helpers import ParamError
 
 
-def g(x: f32[" b"]):
+def g(x: Float32[jnp.ndarray, " b"]):
     pass
 
 
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_beartype_full.py

@@ -0,0 +1,62 @@
+# 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 dataclasses
+
+import equinox as eqx
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Float32
+
+from .helpers import ParamError
+
+
+def g(x: Float32[jnp.ndarray, " b"]):
+    pass
+
+
+g(jnp.array([1.0]))
+with pytest.raises(ParamError):
+    g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_beartype_old.py

@@ -0,0 +1,62 @@
+# 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 dataclasses
+
+import equinox as eqx
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Float32
+
+from .helpers import ParamError
+
+
+def g(x: Float32[jnp.ndarray, " b"]):
+    pass
+
+
+g(jnp.array([1.0]))
+with pytest.raises(ParamError):
+    g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_broken_checker.py

@@ -17,17 +17,46 @@
 # 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 dataclasses
+
+import equinox as eqx
 import jax.numpy as jnp
 import pytest
-from helpers import ParamError
 
-from jaxtyping import f32
+from jaxtyping import Float32
+
+from .helpers import ParamError
 
 
-def g(x: f32[" b"]):
+def g(x: Float32[jnp.ndarray, " b"]):
     pass
 
 
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_transitive/__init__.py

@@ -17,4 +17,4 @@
 # 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
+from . import another_file  # noqa: F401

test/import_hook_tester_transitive/another_file.py

@@ -19,15 +19,30 @@
 
 import jax.numpy as jnp
 import pytest
-from helpers import ParamError
 
-from jaxtyping import f32
+from jaxtyping import Float32
+
+from ..helpers import ParamError
 
 
-def g(x: f32[" b"]):
+def g(x: Float32[jnp.ndarray, " b"]):
     pass
 
 
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+# Typeguard 3.0 no longer supports this
+#
+# class M(eqx.Module):
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+
+
+# M(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1, jnp.array(1.0))

test/import_hook_tester_typeguard.py

@@ -17,17 +17,47 @@
 # 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
+from jaxtyping import Float32
+
+from .helpers import ParamError
 
 
-def g(x: f32[" b"]):
+def g(x: Float32[jnp.ndarray, " b"]):
     pass
 
 
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+# Typeguard 3.0 no longer supports this.
+#
+# class M(eqx.Module):
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# M(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1, jnp.array(1.0))
+#
+#
+#
+# @dataclasses.dataclass
+# class D:
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# D(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1, jnp.array(1.0))

test/import_hook_tester_typeguard_old.py

@@ -0,0 +1,63 @@
+# 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 jaxtyping import Float32
+
+from .helpers import ParamError
+
+
+def g(x: Float32[jnp.ndarray, " b"]):
+    pass
+
+
+g(jnp.array([1.0]))
+with pytest.raises(ParamError):
+    g(jnp.array(1))
+
+
+# Typeguard 3.0 no longer supports this.
+#
+# class M(eqx.Module):
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# M(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1, jnp.array(1.0))
+#
+#
+#
+# @dataclasses.dataclass
+# class D:
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# D(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1, jnp.array(1.0))

test/requirements.txt

@@ -2,3 +2,4 @@ equinox>=0.5.3
 pytest>=7.0.1
 beartype>=0.10.4
 typeguard>=2.13.3
+cloudpickle>=2.2.1

test/test_array.py

@@ -17,34 +17,81 @@
 # 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 typing import get_args, get_origin, Union
+
 import jax.numpy as jnp
 import jax.random as jr
+import numpy as np
 import pytest
-from helpers import ParamError, ReturnError
+import torch
 
-from jaxtyping import Array, f, f32, jaxtyped
+from jaxtyping import (
+    AbstractDtype,
+    Array,
+    ArrayLike,
+    Float,
+    Float32,
+    jaxtyped,
+    PRNGKeyArray,
+    Shaped,
+)
+
+from .helpers import ParamError, ReturnError
 
 
 def test_basic(typecheck):
     @jaxtyped
     @typecheck
-    def g(x: Array["..."]):
+    def g(x: Shaped[Array, "..."]):
         pass
 
     g(jnp.array(1.0))
 
 
+def test_dtypes():
+    from jaxtyping import (  # noqa: F401
+        Array,
+        BFloat16,
+        Bool,
+        Complex,
+        Complex64,
+        Complex128,
+        Float,
+        Float16,
+        Float32,
+        Float64,
+        Inexact,
+        Int,
+        Int8,
+        Int16,
+        Int32,
+        Int64,
+        Num,
+        Shaped,
+        UInt,
+        UInt8,
+        UInt16,
+        UInt32,
+        UInt64,
+    )
+
+    for key, val in locals().items():
+        if issubclass(val, AbstractDtype):
+            assert key == val.__name__
+
+
 def test_return(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f["b c"]) -> f["c b"]:
+    def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]:
         return jnp.transpose(x)
 
     g(jr.normal(getkey(), (3, 4)))
 
     @jaxtyped
     @typecheck
-    def h(x: f["b c"]) -> f["b c"]:
+    def h(x: Float[Array, "b c"]) -> Float[Array, "b c"]:
         return jnp.transpose(x)
 
     with pytest.raises(ReturnError):
@@ -54,7 +101,7 @@ def test_return(typecheck, getkey):
 def test_two_args(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: Array["b c"], y: Array["c d"]):
+    def g(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]):
         return x @ y
 
     g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (4, 5)))
@@ -63,7 +110,7 @@ def test_two_args(typecheck, getkey):
 
     @jaxtyped
     @typecheck
-    def h(x: Array["b c"], y: Array["c d"]) -> Array["b d"]:
+    def h(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]) -> Shaped[Array, "b d"]:
         return x @ y
 
     h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (4, 5)))
@@ -74,7 +121,7 @@ def test_two_args(typecheck, getkey):
 def test_any_dtype(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: Array["a b"]) -> Array["a b"]:
+    def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
         return x
 
     g(jr.normal(getkey(), (3, 4)))
@@ -91,12 +138,12 @@ def test_any_dtype(typecheck, getkey):
 def test_nested_jaxtyped(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32["b c"], transpose: bool) -> f32["c b"]:
+    def g(x: Float32[Array, "b c"], transpose: bool) -> Float32[Array, "c b"]:
         return h(x, transpose)
 
     @jaxtyped
     @typecheck
-    def h(x: f32["c b"], transpose: bool) -> f32["b c"]:
+    def h(x: Float32[Array, "c b"], transpose: bool) -> Float32[Array, "b c"]:
         if transpose:
             return jnp.transpose(x)
         else:
@@ -112,11 +159,11 @@ def test_nested_jaxtyped(typecheck, getkey):
 def test_nested_nojaxtyped(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32["b c"]):
+    def g(x: Float32[Array, "b c"]):
         return h(x)
 
     @typecheck
-    def h(x: f32["c b"]):
+    def h(x: Float32[Array, "c b"]):
         return x
 
     with pytest.raises(ParamError):
@@ -126,14 +173,14 @@ def test_nested_nojaxtyped(typecheck, getkey):
 def test_isinstance(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32["b c"]) -> f32[" z"]:
+    def g(x: Float32[Array, "b c"]) -> Float32[Array, " z"]:
         y = jnp.transpose(x)
-        assert isinstance(y, f32["c b"])
+        assert isinstance(y, Float32[Array, "c b"])
         assert not isinstance(
-            y, f32["b z"]
+            y, Float32[Array, "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
+        assert isinstance(out, Float32[Array, "z"])  # z now bound
         return out
 
     g(jr.normal(getkey(), (2, 3)))
@@ -142,7 +189,9 @@ def test_isinstance(typecheck, getkey):
 def test_fixed(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32["4 5 foo"], y: f32[" foo"]) -> f32["4 5"]:
+    def g(
+        x: Float32[Array, "4 5 foo"], y: Float32[Array, " foo"]
+    ) -> Float32[Array, "4 5"]:
         return x @ y
 
     a = jr.normal(getkey(), (4, 5, 2))
@@ -157,7 +206,7 @@ def test_fixed(typecheck, getkey):
 def test_anonymous(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32["foo _"], y: f32[" _"]):
+    def g(x: Float32[Array, "foo _"], y: Float32[Array, " _"]):
         pass
 
     a = jr.normal(getkey(), (3, 4))
@@ -168,7 +217,11 @@ def test_anonymous(typecheck, getkey):
 def test_named_variadic(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32["*batch foo"], y: f32[" *batch"], z: f32[" foo"]):
+    def g(
+        x: Float32[Array, "*batch foo"],
+        y: Float32[Array, " *batch"],
+        z: Float32[Array, " foo"],
+    ):
         pass
 
     c = jr.normal(getkey(), (5,))
@@ -188,7 +241,7 @@ def test_named_variadic(typecheck, getkey):
 
     @jaxtyped
     @typecheck
-    def h(x: f32[" foo *batch"], y: f32[" foo *batch bar"]):
+    def h(x: Float32[Array, " foo *batch"], y: Float32[Array, " foo *batch bar"]):
         pass
 
     a = jr.normal(getkey(), (4,))
@@ -204,7 +257,7 @@ def test_named_variadic(typecheck, getkey):
 def test_anonymous_variadic(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32["... foo"], y: f32[" foo"]):
+    def g(x: Float32[Array, "... foo"], y: Float32[Array, " foo"]):
         pass
 
     a1 = jr.normal(getkey(), (5,))
@@ -226,7 +279,7 @@ def test_anonymous_variadic(typecheck, getkey):
 def test_broadcast_fixed(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32["4#"]):
+    def g(x: Float32[Array, "#4"]):
         pass
 
     g(jr.normal(getkey(), (4,)))
@@ -239,7 +292,7 @@ def test_broadcast_fixed(typecheck, getkey):
 def test_broadcast_named(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32[" foo#"], y: f32[" foo#"]):
+    def g(x: Float32[Array, " #foo"], y: Float32[Array, " #foo"]):
         pass
 
     a = jr.normal(getkey(), (3,))
@@ -263,7 +316,7 @@ def test_broadcast_named(typecheck, getkey):
 def test_broadcast_variadic_named(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: f32[" *foo#"], y: f32[" *foo#"]):
+    def g(x: Float32[Array, " *#foo"], y: Float32[Array, " *#foo"]):
         pass
 
     a = jr.normal(getkey(), (3,))
@@ -282,12 +335,11 @@ def test_broadcast_variadic_named(typecheck, getkey):
     g(b, b)
     g(c, c)
     g(d, d)
+    g(b, c)
     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):
@@ -295,26 +347,20 @@ def test_broadcast_variadic_named(typecheck, getkey):
 
     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, j)
+    g(d, j)
+    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(a, m)
     g(c, m)
     g(d, m)
-    with pytest.raises(ParamError):
-        g(a, n)
-    with pytest.raises(ParamError):
-        g(b, n)
+    g(a, n)
+    g(b, n)
     with pytest.raises(ParamError):
         g(c, n)
     with pytest.raises(ParamError):
@@ -326,6 +372,206 @@ def test_broadcast_variadic_named(typecheck, getkey):
         g(o, a)
 
 
-def test_no_commas(typecheck, getkey):
+def test_no_commas():
     with pytest.raises(ValueError):
-        f32["foo, bar"]
+        Float32[Array, "foo, bar"]
+
+
+def test_symbolic(typecheck, getkey):
+    @jaxtyped
+    @typecheck
+    def make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
+        return x[1:]
+
+    @jaxtyped
+    @typecheck
+    def cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
+        return jnp.concatenate([x, x])
+
+    @jaxtyped
+    @typecheck
+    def bad_make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
+        return x
+
+    @jaxtyped
+    @typecheck
+    def bad_cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
+        return jnp.concatenate([x, x, x])
+
+    x = jr.normal(getkey(), (5,))
+    assert make_slice(x).shape == (4,)
+    assert cat(x).shape == (10,)
+
+    y = jr.normal(getkey(), (3, 4))
+    with pytest.raises(ParamError):
+        make_slice(y)
+    with pytest.raises(ParamError):
+        cat(y)
+
+    with pytest.raises(ReturnError):
+        bad_make_slice(x)
+    with pytest.raises(ReturnError):
+        bad_cat(x)
+
+
+def test_incomplete_symbolic(typecheck, getkey):
+    @jaxtyped
+    @typecheck
+    def foo(x: Float32[Array, " 2*dim"]):
+        pass
+
+    x = jr.normal(getkey(), (4,))
+    with pytest.raises(NameError):
+        foo(x)
+
+
+def test_arraylike(typecheck, getkey):
+    floatlike1 = Float32[ArrayLike, ""]
+    floatlike2 = Float[ArrayLike, ""]
+    floatlike3 = Float32[ArrayLike, "4"]
+
+    assert get_origin(floatlike1) is Union
+    assert get_origin(floatlike2) is Union
+    assert get_origin(floatlike3) is Union
+    assert set(get_args(floatlike1)) == {
+        Float32[Array, ""],
+        Float32[np.ndarray, ""],
+        Float32[np.bool_, ""],
+        Float32[np.number, ""],
+        float,
+    }
+    assert set(get_args(floatlike2)) == {
+        Float[Array, ""],
+        Float[np.ndarray, ""],
+        Float[np.bool_, ""],
+        Float[np.number, ""],
+        float,
+    }
+    assert set(get_args(floatlike3)) == {
+        Float32[Array, "4"],
+        Float32[np.ndarray, "4"],
+        Float32[np.bool_, "4"],
+        Float32[np.number, "4"],
+    }
+
+    shaped1 = Shaped[ArrayLike, ""]
+    shaped2 = Shaped[ArrayLike, "4"]
+    assert get_origin(shaped1) is Union
+    assert get_origin(shaped2) is Union
+    assert set(get_args(shaped1)) == {
+        Shaped[Array, ""],
+        Shaped[np.ndarray, ""],
+        Shaped[np.bool_, ""],
+        Shaped[np.number, ""],
+        bool,
+        int,
+        float,
+        complex,
+    }
+    assert set(get_args(shaped2)) == {
+        Shaped[Array, "4"],
+        Shaped[np.ndarray, "4"],
+        Shaped[np.bool_, "4"],
+        Shaped[np.number, "4"],
+    }
+
+
+def test_subclass():
+    assert issubclass(Float[Array, ""], Array)
+    assert issubclass(Float[np.ndarray, ""], np.ndarray)
+    assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
+
+
+def test_ignored_names():
+    x = Float[np.ndarray, "foo=4"]
+
+    assert isinstance(np.zeros(4), x)
+    assert not isinstance(np.zeros(5), x)
+    assert not isinstance(np.zeros((4, 5)), x)
+
+    y = Float[np.ndarray, "bar qux foo=bar+qux"]
+
+    assert isinstance(np.zeros((2, 3, 5)), y)
+    assert not isinstance(np.zeros((2, 3, 6)), y)
+
+    z = Float[np.ndarray, "bar #foo=bar"]
+
+    assert isinstance(np.zeros((3, 3)), z)
+    assert isinstance(np.zeros((3, 1)), z)
+    assert not isinstance(np.zeros((3, 4)), z)
+
+    # Weird but legal
+    w = Float[np.ndarray, "bar foo=#bar"]
+
+    assert isinstance(np.zeros((3, 3)), w)
+    assert isinstance(np.zeros((3, 1)), w)
+    assert not isinstance(np.zeros((3, 4)), w)
+
+
+def test_symbolic_functions():
+    x = Float[np.ndarray, "foo bar min(foo,bar)"]
+
+    assert isinstance(np.zeros((2, 3, 2)), x)
+    assert isinstance(np.zeros((3, 2, 2)), x)
+    assert not isinstance(np.zeros((3, 2, 4)), x)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10")
+def test_py310_unions():
+    x = np.zeros(3)
+    y = Shaped[Array | np.ndarray, "_"]
+    assert isinstance(x, get_args(y))
+
+
+def test_key(typecheck):
+    @jaxtyped
+    @typecheck
+    def f(x: PRNGKeyArray):
+        pass
+
+    x = jr.PRNGKey(0)
+    f(x)
+
+    with pytest.raises(ParamError):
+        f(object())
+    with pytest.raises(ParamError):
+        f(1)
+    with pytest.raises(ParamError):
+        f(jnp.array(3))
+    with pytest.raises(ParamError):
+        f(jnp.array(3.0))
+
+
+def test_extension(typecheck, getkey):
+    X = Shaped[Array, "a b"]
+    Y = Shaped[X, "c d"]
+    Z = Shaped[Array, "c d a b"]
+    assert str(Z) == str(Y)
+
+    X = Float[Array, "a"]
+    Y = Float[X, "b"]
+
+    @jaxtyped
+    @typecheck
+    def f(a: X, b: Y):
+        ...
+
+    a = jr.normal(getkey(), (3, 4))
+    b = jr.normal(getkey(), (4,))
+    c = jr.normal(getkey(), (3,))
+
+    f(b, a)
+    with pytest.raises(ParamError):
+        f(c, a)
+    with pytest.raises(ParamError):
+        f(a, a)
+
+    @typecheck
+    def g(a: Shaped[PRNGKeyArray, "2"]):
+        ...
+
+    with pytest.raises(ParamError):
+        g(jr.PRNGKey(0))
+    g(jr.split(jr.PRNGKey(0)))
+    with pytest.raises(ParamError):
+        g(jr.split(jr.PRNGKey(0), 3))

test/test_decorator.py

@@ -0,0 +1,77 @@
+import abc
+
+from jaxtyping import jaxtyped
+
+
+class M(metaclass=abc.ABCMeta):
+    @jaxtyped
+    def f(self):
+        ...
+
+    @jaxtyped
+    @classmethod
+    def g1(cls):
+        return 3
+
+    @classmethod
+    @jaxtyped
+    def g2(cls):
+        return 4
+
+    @jaxtyped
+    @staticmethod
+    def h1():
+        return 3
+
+    @staticmethod
+    @jaxtyped
+    def h2():
+        return 4
+
+    @jaxtyped
+    @abc.abstractmethod
+    def i1(self):
+        ...
+
+    @abc.abstractmethod
+    @jaxtyped
+    def i2(self):
+        ...
+
+
+class N:
+    @jaxtyped
+    @property
+    def j1(self):
+        return 3
+
+    @property
+    @jaxtyped
+    def j2(self):
+        return 4
+
+
+def test_identity():
+    assert M.f is M.f
+
+
+def test_classmethod():
+    assert M.g1() == 3
+    assert M.g2() == 4
+
+
+def test_staticmethod():
+    assert M.h1() == 3
+    assert M.h2() == 4
+
+
+# Check that the @jaxtyped decorator doesn't blat the __isabstractmethod__ of
+# @abstractmethod
+def test_abstractmethod():
+    assert M.i1.__isabstractmethod__
+    assert M.i2.__isabstractmethod__
+
+
+def test_property():
+    assert N().j1 == 3
+    assert N().j2 == 4

test/test_import_hook.py

@@ -22,35 +22,71 @@ import pytest
 from jaxtyping import install_import_hook
 
 
+def test_import_hook_typeguard_old():
+    hook = install_import_hook(
+        "test.import_hook_tester_typeguard_old", ("typeguard", "typechecked")
+    )
+    with hook:
+        from . import import_hook_tester_typeguard_old  # noqa: F401
+
+
 def test_import_hook_typeguard():
     hook = install_import_hook(
-        "import_hook_tester_typeguard", ("typeguard", "typechecked")
+        "test.import_hook_tester_typeguard", "typeguard.typechecked"
     )
-    import import_hook_tester_typeguard  # noqa: F401
+    with hook:
+        from . import import_hook_tester_typeguard  # noqa: F401
 
-    hook.uninstall()
+
+def test_import_hook_beartype_old():
+    try:
+        import beartype  # noqa: F401
+    except ImportError:
+        pytest.skip("Beartype not installed")
+    else:
+        hook = install_import_hook(
+            "test.import_hook_tester_beartype_old", ("beartype", "beartype")
+        )
+        with hook:
+            from . import import_hook_tester_beartype_old  # noqa: F401
 
 
 def test_import_hook_beartype():
-    hook = install_import_hook("import_hook_tester_beartype", ("beartype", "beartype"))
-    import import_hook_tester_beartype  # noqa: F401
+    try:
+        import beartype  # noqa: F401
+    except ImportError:
+        pytest.skip("Beartype not installed")
+    else:
+        hook = install_import_hook(
+            "test.import_hook_tester_beartype", "beartype.beartype"
+        )
+        with hook:
+            from . import import_hook_tester_beartype  # noqa: F401
 
-    hook.uninstall()
+
+def test_import_hook_beartype_full():
+    try:
+        import beartype  # noqa: F401
+    except ImportError:
+        pytest.skip("Beartype not installed")
+    else:
+        bearchecker = "beartype.beartype(conf=beartype.BeartypeConf(strategy=beartype.BeartypeStrategy.On))"  # noqa: E501
+        hook = install_import_hook("test.import_hook_tester_beartype_full", bearchecker)
+        with hook:
+            from . import import_hook_tester_beartype_full  # noqa: F401
 
 
 def test_import_hook_transitive():
     hook = install_import_hook(
-        "import_hook_tester_transitive", ("typeguard", "typechecked")
+        "test.import_hook_tester_transitive", "typeguard.typechecked"
     )
-    import import_hook_tester_transitive  # noqa: F401
-
-    hook.uninstall()
+    with hook:
+        from . import import_hook_tester_transitive  # noqa: F401
 
 
 def test_import_hook_broken_checker():
     hook = install_import_hook(
-        "import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist")
+        "test.import_hook_tester_broken_checker", "jaxtyping.does_not_exist"
     )
-    with pytest.raises(AttributeError):
-        import import_hook_tester_broken_checker  # noqa: F401
-    hook.uninstall()
+    with hook, pytest.raises(AttributeError):
+        from . import import_hook_tester_broken_checker  # noqa: F401

test/test_pytree.py

@@ -17,16 +17,17 @@
 # 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
+from typing import NamedTuple, 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
+from jaxtyping import Float, jaxtyped, PyTree
+
+from .helpers import make_mlp, ParamError
 
 
 def test_direct(typecheck):
@@ -94,7 +95,7 @@ def test_nested_pytrees(getkey, typecheck):
 def test_pytree_array(typecheck):
     @jaxtyped
     @typecheck
-    def g(x: PyTree[f["..."]]):
+    def g(x: PyTree[Float[jnp.ndarray, "..."]]):
         pass
 
     g(jnp.array(1.0))
@@ -108,7 +109,7 @@ def test_pytree_array(typecheck):
 def test_pytree_shaped_array(typecheck, getkey):
     @jaxtyped
     @typecheck
-    def g(x: PyTree[f["b c"]]):
+    def g(x: PyTree[Float[jnp.ndarray, "b c"]]):
         pass
 
     g(jnp.array([[1.0]]))
@@ -154,3 +155,33 @@ def test_pytree_tuple(typecheck):
         g([1, 1])
     with pytest.raises(ParamError):
         g([(1, 1), "hi"])
+
+
+def test_pytree_namedtuple(typecheck):
+    class CustomNamedTuple(NamedTuple):
+        x: Float[jnp.ndarray, "a b"]
+        y: Float[jnp.ndarray, "b c"]
+
+    class OtherCustomNamedTuple(NamedTuple):
+        x: Float[jnp.ndarray, "a b"]
+        y: Float[jnp.ndarray, "b c"]
+
+    @typecheck
+    def g(x: PyTree[CustomNamedTuple]):
+        ...
+
+    g(
+        CustomNamedTuple(
+            x=jax.random.normal(jax.random.PRNGKey(42), (3, 2)),
+            y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)),
+        )
+    )
+    with pytest.raises(ParamError):
+        g(object())
+    with pytest.raises(ParamError):
+        g(
+            OtherCustomNamedTuple(
+                x=jax.random.normal(jax.random.PRNGKey(42), (3, 2)),
+                y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)),
+            )
+        )

test/test_serialisation.py

@@ -0,0 +1,16 @@
+import cloudpickle
+import numpy as np
+import torch
+
+from jaxtyping import AbstractArray, Array, Shaped
+
+
+def test_pickle():
+    x = cloudpickle.dumps(Shaped[Array, ""])
+    y = cloudpickle.dumps(AbstractArray)
+    z = cloudpickle.dumps(Shaped[np.ndarray, ""])
+    w = cloudpickle.dumps(Shaped[torch.Tensor, ""])
+    cloudpickle.loads(x)
+    cloudpickle.loads(y)
+    cloudpickle.loads(z)
+    cloudpickle.loads(w)

test/test_threading.py

@@ -0,0 +1,64 @@
+# 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 threading
+
+import jax.numpy as jnp
+from typeguard import typechecked
+
+from jaxtyping import Array, Float, jaxtyped
+
+
+class _ErrorableThread(threading.Thread):
+    def run(self):
+        try:
+            super().run()
+        except Exception as e:
+            self.exc = e
+
+    def join(self, timeout=None):
+        super().join(timeout)
+        if hasattr(self, "exc"):
+            raise self.exc
+
+
+def test_threading_jaxtyped():
+    @jaxtyped
+    @typechecked
+    def add(x: Float[Array, "a b"], y: Float[Array, "a b"]) -> Float[Array, "a b"]:
+        return x + y
+
+    def run():
+        a = jnp.array([[1.0, 2.0]])
+        b = jnp.array([[2.0, 3.0]])
+        add(a, b)
+
+    thread = _ErrorableThread(target=run)
+    thread.start()
+    thread.join()
+
+
+def test_threading_nojaxtyped():
+    def run():
+        a = jnp.array([[1.0, 2.0]])
+        assert isinstance(a, Float[Array, "..."])
+
+    thread = _ErrorableThread(target=run)
+    thread.start()
+    thread.join()