version bump

Phew, this ended up being a pretty complicated change!
The basic summary is that we now support the syntax
```
@jaxtyped(typechecker=typechecker)
def f(...): ...
```
and when using this, we now get pretty error messages about what went
wrong.

(
The old syntax, i.e.
```
@jaxtyped
@typechecker
def f(...): ...
```
is still supported, but doesn't give much information.
)

The internals of this do quite a lot of magic! In particular we
dynamically create quite a lot of functions and test the provided
arguments against their signatures. The overhead should still be
minimal under `jax.jit`, though.
(TODO: what's the overhead like in non-jit situations, e.g. PyTorch?
I've tried to minimise the overhead throughout just to be sure, but
perhaps PyTorch users should stick to the old syntax?)

.flake8

@@ -1,4 +0,0 @@
-[flake8]
-max-line-length = 88
-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 -r ${{ github.workspace }}/test/requirements.txt
+                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 -r test/requirements.txt
+          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,7 +0,0 @@
-[settings]
-force_alphabetical_sort_within_sections=true
-lines_after_imports=2
-profile=black
-combine_as_imports=True
-treat_comments_as_code=true
-extra_standard_library=typing_extensions

.pre-commit-config.yaml

@@ -19,20 +19,11 @@
 
 repos:
   - repo: https://github.com/ambv/black
-    rev: 22.3.0
+    rev: 23.9.1
     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.291'
     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
+        args: ["--fix"]

API.md

@@ -1,231 +0,0 @@
-# Full API
-
-## Annotating array types
-
-Each array is denoted by a type `dtype[array, shape]`, such as `Float[Array, "batch channels"]`.
-
-### Shape
-
-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) below.)
-
-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 _ _"`.
-
-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"`
-
-Some 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 (imported from `jaxtyping`):
-- Any dtype at all: `Shaped`
-  - Boolean: `Bool`
-  - 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 typically be a `jaxtyping.Array`, which is an alias for `jax.numpy.ndarray`.
-
-But you can use other types as well. `jaxtyping` has support for JAX, NumPy, TensorFlow, and PyTorch, e.g.:
-```python
-Float[np.ndarray, "..."]
-Float[tf.Tensor, "..."]
-Float[torch.Tensor, "..."]
-```
-
-## PyTrees
-
-### `jaxtyping.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))
-
-## 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.
-
-### Option 1: `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 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, :]
-```
-
-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 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"])`.
-
-### Option 2: `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 hook.
-
-This can be used via a `with` block; for example:
-```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 name begins
-    import foo.bar      # with the specified string, will automatically have both `@jaxtyped`
-    import foo.bar.qux  # 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.
-
-**Example: writing an end-user script**
-
-```python
-### entry_point.py
-from jaxtyping import install_import_hook
-with install_import_hook("do_stuff", ("typeguard", "typechecked")):
-    import do_stuff
-
-### do_stuff.py
-from jaxtyping import Array, Float32
-
-def g(x: Float32[Array, "..."]):
-    ...
-```
-
-**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.
-```
-
-#### 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`.
-
-## Static type checking
-
-jaxtyping should be compatible with static type checkers (the big three are `mypy`, `pyright`, `pytype`) out of the box.
-
-Due to limitations of static type checkers, only the array type (JAX array vs NumPy array vs PyTorch tensor vs TensorFlow tensor) is checked. Shape and dtype are not checked. [See the FAQ](./FAQ.md) for more details.
-
-## 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 `Float`, `Inexact` etc.) For example:
-```python
-class UInt8or16(AbstractDtype):
-    dtypes = ["uint8", "uint16"]
-
-UInt8or16[Array, "shape"]
-```
-which is functionally equivalent to
-```python
-Union[UInt8[Array, "shape"], UInt16[Array, "shape"]]
-```
-
-### `jaxtyping.AbstractArray`
-
-The base class of all shape-and-dtype-specified arrays, e.g. it's a base class
-for `Float32[Array, "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
 ```
 

MANIFEST.in

@@ -1,2 +0,0 @@
-include jaxtyping/py.typed
-prune test

README.md

@@ -1,15 +1,16 @@
 <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 Array, Float, PyTree
 
-# Accepts floating-point 2D arrays with matching dimensions
+# Accepts floating-point 2D arrays with matching axes
 def matrix_multiply(x: Float[Array, "dim1 dim2"],
                     y: Float[Array, "dim2 dim3"]
                   ) -> Float[Array, "dim1 dim3"]:
@@ -28,31 +29,41 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 pip install jaxtyping
 ```
 
-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).
+Requires Python 3.9+.
+
+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
 
-### See also: other tools in the JAX ecosystem
+### See also: other libraries in the JAX ecosystem
 
-Neural networks: [Equinox](https://github.com/patrick-kidger/equinox).
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
 
-Numerical differential equation solvers: [Diffrax](https://github.com/patrick-kidger/diffrax).
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
 
-Computer vision models: [Eqxvision](https://github.com/paganpasta/eqxvision).
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
 
-SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2jax](https://github.com/google/sympy2jax).
+[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.
 
-### Acknowledgements
+[Lineax](https://github.com/google/lineax): linear solvers.
 
-Shape annotations + runtime type checking is inspired by [TorchTyping](https://github.com/patrick-kidger/torchtyping).
+[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.
 
-The concise syntax is partially inspired by [etils.array_types](https://github.com/google/etils/tree/main/etils/array_types).
+[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).
+
+[sympy2jax](https://github.com/google/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.
+
+[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
+
+[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).
+
+[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)
 
 ### Disclaimer
 

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,18 @@
+# Advanced features
+
+## Creating your own dtypes
+
+::: jaxtyping.AbstractDtype
+    selection:
+        members:
+            false
+
+## Introspection
+
+If you're writing your own type hint parser, then you may wish to detect if some Python object is a jaxtyping-provided type.
+
+You can check for dtypes by doing `issubclass(x, AbstractDtype)`. For example, `issubclass(Float32, AbstractDtype)` will pass.
+
+You can check for arrays by doing `issubclass(x, AbstractArray)`. Here, `AbstractArray` is the base class for all shape-and-dtype specified arrays, e.g. it's a base class for `Float32[Array, "foo"]`.
+
+You can check for pytrees by doing `issubclass(x, PyTree)`. For example, `issubclass(PyTree[int], PyTree)` will pass.

docs/api/array.md

@@ -0,0 +1,123 @@
+# 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 in terms of other variable-size axes, e.g.  
+    `def remove_last(x: Float[Array, "dim"]) -> Float[Array, "dim-1"]`.  
+    Symbolic expressions must not use any spaces, otherwise each piece is treated as as a separate axis.
+
+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 an axis to indicate that it can match multiple axes, e.g. `"*batch"` will match zero or more batch axes.
+- Prepend `#` to an axis 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 an axis to disable any runtime checking of that axis (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"]`.
+- Prepend `?` to an axis to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree/).)
+
+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.
+- 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.
+- Symbolic expressions are evaluated in two stages: they are first evaluated as f-strings using the arguments of the function, and second are evaluated using the processed axis sizes. The f-string evaluation means that they can use local variables by enclosing them with curly braces, e.g. `{variable}`, e.g.
+    ```python
+    def full(size: int, fill: float) -> Float[Array, "{size}"]:
+        return jax.numpy.full((size,), fill)
+
+    class SomeClass:
+        some_value = 5
+
+        def full(self, fill: float) -> Float[Array, "{self.some_value}+3"]:
+            return jax.numpy.full((self.some_value + 3,), fill)
+    ```
+
+## 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`
+        - Any floating, integer, or unsigned integer: `Real`.
+
+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, PRNG keys
+
+For convenience, jaxtyping also includes `jaxtyping.Scalar`, `jaxtyping.ScalarLike`, and `jaxtyping.PRNGKeyArray`, defined as:
+```python
+Scalar = Shaped[Array, ""]
+ScalarLike = Shaped[ArrayLike, ""]
+
+# Left: new-style typed keys; right: old-style keys. See JEP 9263.
+PRNGKeyArray = Union[Key[Array, ""], 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,46 @@
+# PyTree annotations
+
+:::jaxtyping.PyTree
+    selection:
+        members:
+            false
+
+---
+
+:::jaxtyping.PyTreeDef
+
+---
+
+## Path-dependent shapes
+
+The prefix `?` may be used to indicate that the axis size can depend on which leaf of a PyTree the array is at. For example:
+```python
+def f(
+    x: PyTree[Shaped[Array, "?foo"], "T"],
+    y: PyTree[Shaped[Array, "?foo"], "T"],
+):
+    pass
+```
+The above demands that `x` and `y` have matching PyTree structures (due to the `T` annotation), and that their leaves must all be one-dimensional arrays, *and that the corresponding pairs of leaves in `x` and `y` must have the same size as each other*.
+
+Thus the following is allowed:
+```python
+x0 = jnp.arange(3)
+x1 = jnp.arange(5)
+
+y0 = jnp.arange(3) + 1
+y1 = jnp.arange(5) + 1
+
+f((x0, x1), (y0, y1))  # x0 matches y0, and x1 matches y1. All good!
+```
+
+But this is not:
+```python
+f((x1, x1), (y0, y1))  # x1 does not have a size matching y0!
+```
+
+Internally, all that is happening is that `foo` is replaced with `0foo` for the first leaf, `1foo` for the next leaf, etc., so that each leaf gets a unique version of the name.
+
+---
+
+Note that `jaxtyping.{PyTree, PyTreeDef}` are only available if JAX has been installed.

docs/api/runtime-type-checking.md

@@ -0,0 +1,51 @@
+# 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 only whilst tracing, and will not impact runtime performance.
+
+There are two approaches: either use [`jaxtyping.jaxtyped`][] to typecheck a single function, or [`jaxtyping.install_import_hook`][] to typecheck a whole codebase.
+
+In either case, the actual business of checking types is performed with the help of a runtime type-checking library. The two most popular are [beartype](https://github.com/beartype/beartype) and [typeguard](https://github.com/agronholm/typeguard). (If using typeguard, then specifically the version `2.*` series should be used. Later versions -- `3` and `4` -- have some known issues.)
+
+---
+
+::: jaxtyping.jaxtyped
+
+---
+
+::: jaxtyping.install_import_hook
+
+---
+
+#### 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`.
+
+#### IPython extension
+
+If you are running in an IPython environment (for example a Jupyter or Colab notebook), then the jaxtyping hook can be automatically ran via a custom magic:
+```python
+import jaxtyping
+%load_ext jaxtyping
+%jaxtyping.typechecker beartype.beartype  # or any other runtime type checker
+```
+Place this at the start of your notebook -- everything that is directly defined in the notebook, after this magic is run, will be hook'd.
+
+#### Other runtime type-checking libraries
+
+Beartype and typeguard happen to be the two most popular runtime type-checking libraries (at least at time of writing), but jaxtyping should be compatible with all runtime type checkers out-of-the-box. The runtime type-checking library just needs to provide a type-checking decorator (analgous to `beartype.beartype` or `typeguard.typechecked`), and perform `isinstance` checks against jaxtyping's types.

docs/faq.md

@@ -1,6 +1,6 @@
 # FAQ
 
-## Does jaxtyping work with static type checkers like `mypy`/`pyright`/`pytype`?
+## 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.
 
@@ -12,13 +12,13 @@ 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` is throwing an error.
+## `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`, at least, this is easily 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.
+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)?
 

docs/index.md

@@ -0,0 +1,58 @@
+# 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.9+.
+
+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 axes
+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!
+
+## See also: other libraries in the JAX ecosystem
+
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
+
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
+
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
+
+[Lineax](https://github.com/google/lineax): linear solvers and linear least squares.
+
+[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
+
+[sympy2jax](https://github.com/google/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.
+
+[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).

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,24 +17,29 @@
 # 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
-import typing_extensions
+import warnings
+
+# First import some things as normal
+from ._array_types import (
+    AbstractArray as AbstractArray,
+    AbstractDtype as AbstractDtype,
+    get_array_name_format as get_array_name_format,
+    has_jax,
+    set_array_name_format as set_array_name_format,
+)
+from ._decorator import jaxtyped as jaxtyped, TypeCheckError as TypeCheckError
+from ._import_hook import install_import_hook as install_import_hook
+from ._ipython_extension import load_ipython_extension as load_ipython_extension
 
 
-try:
-    import jax
-except ImportError:
-    has_jax = False
-else:
-    has_jax = True
-    del jax
-
-
-# Type checkers don't know which branch below will be executed.
+# 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):
 
@@ -42,51 +47,181 @@ elif has_jax:
             pass
 
         Array.__module__ = "builtins"
+
+        class ArrayLike:
+            pass
+
+        ArrayLike.__module__ = "builtins"
     else:
         from jax import Array as Array
 
-from .array_types import (
-    AbstractArray as AbstractArray,
-    AbstractDtype as AbstractDtype,
-    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,
-    get_array_name_format as get_array_name_format,
-    Inexact as Inexact,
-    Int as Int,
-    Int8 as Int8,
-    Int16 as Int16,
-    Int32 as Int32,
-    Int64 as Int64,
-    Integer as Integer,
-    Num as Num,
-    set_array_name_format as set_array_name_format,
-    Shaped as Shaped,
-    UInt as UInt,
-    UInt8 as UInt8,
-    UInt16 as UInt16,
-    UInt32 as UInt32,
-    UInt64 as UInt64,
-)
-from .decorator import jaxtyped as jaxtyped
-from .import_hook import install_import_hook as install_import_hook
+        try:
+            from jax.typing import ArrayLike as ArrayLike
+        except (ModuleNotFoundError, ImportError):
+            pass
 
 
+# Import our dtypes
 if typing.TYPE_CHECKING:
-    _T = typing.TypeVar("_T")
+    # Introduce an indirection so that we can `import X as X` to make it clear that
+    # these are public.
+    from jax.typing import DTypeLike as DTypeLike
 
-    class PyTree(typing_extensions.Protocol[_T]):
-        pass
+    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,
+        Real as Real,
+        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,
+        Num as Num,
+        Real as Real,
+        Shaped as Shaped,
+        UInt as UInt,
+        UInt8 as UInt8,
+        UInt16 as UInt16,
+        UInt32 as UInt32,
+        UInt64 as UInt64,
+    )
 
+    if has_jax:
+        import jax.typing
+
+        from ._array_types import Key as Key
+
+        if hasattr(jax.typing, "DTypeLike"):
+            from jax.typing import DTypeLike as DTypeLike
+
+
+# Now import PyTreeDef and PyTree
+if typing.TYPE_CHECKING:
+    import typing_extensions
+
+    from jax.tree_util import PyTreeDef as PyTreeDef
+
+    # Set up to deliberately confuse a static type checker.
+    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
+    if hasattr(typing, "GENERATING_DOCUMENTATION"):
+        # Most parts of the Equinox ecosystem have
+        # `typing.GENERATING_DOCUMENTATION = True` when generating documentation, to
+        # add whatever shims are necessary to get pretty docs. E.g. to have type
+        # annotations appear as just `PyTree`, not `jaxtyping.PyTree`.
+        #
+        # As jaxtyping actually wants things to appear as e.g. `jaxtyping.PyTree`,
+        # rather than just `PyTree`, then it sets
+        # `typing.GENERATING_DOCUMENTATION = False`, to disable these shims.
+        #
+        # Here we do only a `hasattr` check, as we want to get this version of
+        # `PyTreeDef` in both the jaxtyping and the Equinox(/etc.) docs.
+
+        class PyTreeDef:
+            """Alias for `jax.tree_util.PyTreeDef`, which is the type of the return
+            from `jax.tree_util.tree_structure(...)`.
+            """
+
+        if typing.GENERATING_DOCUMENTATION:
+            # Equinox etc. docs get just `PyTreeDef`.
+            # jaxtyping docs get `jaxtyping.PyTreeDef`.
+            PyTreeDef.__module__ = "builtins"
+
+    else:
+        from jax.tree_util import PyTreeDef as PyTreeDef
+
+    from ._pytree_type import PyTree as PyTree  # noqa: F401
+
+
+# Conveniences
+if typing.TYPE_CHECKING:
+    from ._indirection import (
+        PRNGKeyArray as PRNGKeyArray,
+        Scalar as Scalar,
+        ScalarLike as ScalarLike,
+    )
+elif has_jax:
+    from ._array_types import Scalar, ScalarLike  # noqa: F401
+
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        # That is, we're generating some downstream documentation, not the jaxtyping
+        # documentation itself.
+        class PRNGKeyArray:
+            pass
+
+        PRNGKeyArray.__module__ = "builtins"
+    else:
+        from ._array_types import PRNGKeyArray
 
 del has_jax
 
-__version__ = "0.2.10"
+
+check_equinox_version = True  # easy-to-replace line with copybara
+if check_equinox_version:
+    try:
+        eqx_version = importlib.metadata.version("equinox")
+    except importlib.metadata.PackageNotFoundError:
+        pass
+    else:
+        major, minor, patch = eqx_version.split(".")
+        equinox_version = (int(major), int(minor), int(patch))
+        if equinox_version < (0, 11, 0):
+            warnings.warn(
+                "jaxtyping version >=0.2.23 should be used with Equinox version "
+                ">=0.11.1"
+            )
+
+
+__version__ = importlib.metadata.version("jaxtyping")

jaxtyping/_array_types.py

@@ -0,0 +1,720 @@
+# 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 enum
+import functools as ft
+import re
+import sys
+import types
+import typing
+from typing import Any, Literal, NoReturn, Optional, Union
+
+import numpy as np
+
+from ._raise import jaxtyping_raise, jaxtyping_raise_from
+from ._storage import (
+    get_shape_memo,
+    get_treeflatten_memo,
+    get_treepath_memo,
+    set_shape_memo,
+)
+
+
+try:
+    import jax
+except (ImportError, RuntimeError, AttributeError):
+    # We catch `RuntimeError` as JAX will throw this if it's present, but unable to run
+    # on the current machine. This fails with this error.
+    # We catch `AttributeError` as the above then leaves the module in a partially
+    # initialised state, which causes subsequent imports to fail with this error.
+    has_jax = False
+else:
+    has_jax = True
+
+
+_array_name_format = "dtype_and_shape"
+
+
+def get_array_name_format():
+    return _array_name_format
+
+
+def set_array_name_format(value):
+    global _array_name_format
+    _array_name_format = value
+
+
+_any_dtype = object()
+
+
+_anonymous_dim = object()
+_anonymous_variadic_dim = object()
+
+
+class _DimType(enum.Enum):
+    named = enum.auto()
+    fixed = enum.auto()
+    symbolic = enum.auto()
+
+
+class _NamedDim:
+    def __init__(self, name, broadcastable, treepath):
+        self.name = name
+        self.broadcastable = broadcastable
+        self.treepath = treepath
+
+
+class _NamedVariadicDim:
+    def __init__(self, name, broadcastable, treepath):
+        self.name = name
+        self.broadcastable = broadcastable
+        self.treepath = treepath
+
+
+class _FixedDim:
+    def __init__(self, size, broadcastable):
+        self.size = size
+        self.broadcastable = broadcastable
+
+
+class _SymbolicDim:
+    def __init__(self, elem, broadcastable):
+        self.elem = elem
+        self.broadcastable = broadcastable
+
+
+_AbstractDimOrVariadicDim = Union[
+    Literal[_anonymous_dim],
+    Literal[_anonymous_variadic_dim],
+    _NamedDim,
+    _NamedVariadicDim,
+    _FixedDim,
+    _SymbolicDim,
+]
+_AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim, _SymbolicDim]
+
+
+def _check_dims(
+    cls_dims: list[_AbstractDim],
+    obj_shape: tuple[int, ...],
+    single_memo: dict[str, int],
+    arg_memo: dict[str, Any],
+) -> 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:
+            pass
+        elif cls_dim.broadcastable and obj_size == 1:
+            pass
+        elif type(cls_dim) is _FixedDim:
+            if cls_dim.size != obj_size:
+                return False
+        elif type(cls_dim) is _SymbolicDim:
+            try:
+                # Support f-string syntax.
+                # https://stackoverflow.com/a/53671539/22545467
+                elem = eval(f"f'{cls_dim.elem}'", arg_memo.copy())
+                # Make a copy to avoid `__builtins__` getting added as a key.
+                eval_size = eval(elem, single_memo.copy())
+            except NameError as e:
+                jaxtyping_raise_from(
+                    NameError(
+                        f"Cannot process symbolic axis '{cls_dim.elem}' as "
+                        "some axis names have not been processed. In practice you "
+                        "should usually only use symbolic axes in annotations "
+                        "for return types, referring only to axes annotated for "
+                        "arguments."
+                    ),
+                    e,
+                )
+            if eval_size != obj_size:
+                return False
+        else:
+            assert type(cls_dim) is _NamedDim
+            if cls_dim.treepath:
+                name = get_treepath_memo() + cls_dim.name
+            else:
+                name = cls_dim.name
+            try:
+                cls_size = single_memo[name]
+            except KeyError:
+                single_memo[name] = obj_size
+            else:
+                if cls_size != obj_size:
+                    return False
+    return True
+
+
+class _MetaAbstractArray(type):
+    def __instancecheck__(cls, obj):
+        if not isinstance(obj, cls.array_type):
+            return False
+        if get_treeflatten_memo():
+            return True
+
+        if 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:
+                jaxtyping_raise(
+                    RuntimeError("Unrecognised array/tensor type to extract dtype from")
+                )
+
+        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
+
+        single_memo, variadic_memo, pytree_memo, arg_memo = get_shape_memo()
+        single_memo_bak = single_memo.copy()
+        variadic_memo_bak = variadic_memo.copy()
+        pytree_memo_bak = pytree_memo.copy()
+        arg_memo_bak = arg_memo.copy()
+        try:
+            check = cls._check_shape(obj, single_memo, variadic_memo, arg_memo)
+        except Exception:
+            set_shape_memo(
+                single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
+            )
+            raise
+        if check:
+            return True
+        else:
+            set_shape_memo(
+                single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
+            )
+            return False
+
+    def _check_shape(
+        cls,
+        obj,
+        single_memo: dict[str, int],
+        variadic_memo: dict[str, tuple[bool, tuple[int, ...]]],
+        arg_memo: dict[str, Any],
+    ):
+        if cls.index_variadic is None:
+            if obj.ndim != len(cls.dims):
+                return False
+            return _check_dims(cls.dims, obj.shape, single_memo, arg_memo)
+        else:
+            if obj.ndim < len(cls.dims) - 1:
+                return False
+            i = cls.index_variadic
+            j = -(len(cls.dims) - i - 1)
+            if j == 0:
+                j = None
+            if not _check_dims(cls.dims[:i], obj.shape[:i], single_memo, arg_memo):
+                return False
+            if j is not None and not _check_dims(
+                cls.dims[j:], obj.shape[j:], single_memo, arg_memo
+            ):
+                return False
+            variadic_dim = cls.dims[i]
+            if variadic_dim is _anonymous_variadic_dim:
+                return True
+            else:
+                assert type(variadic_dim) is _NamedVariadicDim
+                if variadic_dim.treepath:
+                    name = get_treepath_memo() + variadic_dim.name
+                else:
+                    name = variadic_dim.name
+                broadcastable = variadic_dim.broadcastable
+                try:
+                    prev_broadcastable, prev_shape = variadic_memo[name]
+                except KeyError:
+                    variadic_memo[name] = (broadcastable, obj.shape[i:j])
+                    return True
+                else:
+                    new_shape = obj.shape[i:j]
+                    if prev_broadcastable:
+                        try:
+                            broadcast_shape = np.broadcast_shapes(new_shape, prev_shape)
+                        except ValueError:  # not broadcastable e.g. (3, 4) and (5,)
+                            return False
+                        if not broadcastable and broadcast_shape != new_shape:
+                            return False
+                        variadic_memo[name] = (broadcastable, broadcast_shape)
+                    else:
+                        if broadcastable:
+                            try:
+                                broadcast_shape = np.broadcast_shapes(
+                                    new_shape, prev_shape
+                                )
+                            except ValueError:  # not broadcastable e.g. (3, 4) and (5,)
+                                return False
+                            if broadcast_shape != prev_shape:
+                                return False
+                        else:
+                            if new_shape != prev_shape:
+                                return False
+                    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):
+    for dim in dims:
+        if dim is not _anonymous_variadic_dim and not isinstance(
+            dim, _NamedVariadicDim
+        ):
+            return False
+    return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes)
+
+
+class AbstractArray(metaclass=_MetaAbstractArray):
+    """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 axes.
+            raise ValueError("Axes should be separated with spaces, not commas")
+        if elem.endswith("#"):
+            raise ValueError(
+                "As of jaxtyping v0.1.0, broadcastable axes are now denoted "
+                "with a # at the start, rather than at the end"
+            )
+
+        if "..." in elem:
+            if elem != "...":
+                raise ValueError(
+                    "Anonymous multiple axes '...' must be used on its own; "
+                    f"got {elem}"
+                )
+            broadcastable = False
+            variadic = True
+            anonymous = True
+            treepath = False
+            dim_type = _DimType.named
+        else:
+            broadcastable = False
+            variadic = False
+            anonymous = False
+            treepath = 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 "
+                            "axes, 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:]
+                elif first_char == "?":
+                    if treepath:
+                        raise ValueError(
+                            "Do not use ? twice to denote dependence on location "
+                            "within a PyTree, e.g. `??foo` is not allowed"
+                        )
+                    treepath = 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 variadic 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 axes, e.g. "
+                    "`*4` is not allowed."
+                )
+            if anonymous:
+                raise ValueError(
+                    "Cannot have a fixed axis be anonymous, e.g. `_4` is not allowed."
+                )
+            if treepath:
+                raise ValueError(
+                    "Cannot have a fixed axis have tree-path dependence, e.g. `?4` is "
+                    "not allowed."
+                )
+            elem = _FixedDim(elem, broadcastable)
+        elif dim_type is _DimType.named:
+            if anonymous:
+                if broadcastable:
+                    raise ValueError(
+                        "Cannot have an axis 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, treepath)
+                else:
+                    elem = _NamedDim(elem, broadcastable, treepath)
+        else:
+            assert dim_type is _DimType.symbolic
+            if anonymous:
+                raise ValueError(
+                    "Cannot have a symbolic axis be anonymous, e.g. "
+                    "`_foo+bar` is not allowed"
+                )
+            if variadic:
+                raise ValueError(
+                    "Cannot have symbolic multiple-axes, e.g. "
+                    "`*foo+bar` is not allowed"
+                )
+            if treepath:
+                raise ValueError(
+                    "Cannot have a symbolic axis with tree-path dependence, e.g. "
+                    "`?foo+bar` is not allowed"
+                )
+            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
+    elif array_type is np.bool_:
+        if _check_scalar("bool", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is np.generic or array_type is np.number:
+        if _check_scalar("", 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 variadic 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:
+        jaxtyping_raise(
+            RuntimeError(
+                f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to "
+                "check just the dtype of an array, then use "
+                f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
+            )
+        )
+
+    def __getitem__(cls, item: tuple[Any, str]):
+        if not isinstance(item, tuple) or len(item) != 2:
+            raise ValueError(
+                "As of jaxtyping v0.2.0, type annotations must now include both an "
+                "array type and a shape. For example `Float[Array, 'foo bar']`.\n"
+                "Ellipsis can be used to accept any shape: `Float[Array, '...']`."
+            )
+        array_type, dim_str = item
+        dim_str = dim_str.strip()
+        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:
+                out = Union[out]
+        else:
+            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):
+    """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. "
+            '`Float32("shape")` when you mean `Float32[jnp.ndarray, "shape"]`?'
+        )
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+
+        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
+
+
+_prng_key = "prng_key"
+_bool = "bool"
+_bool_ = "bool_"
+_uint8 = "uint8"
+_uint16 = "uint16"
+_uint32 = "uint32"
+_uint64 = "uint64"
+_int8 = "int8"
+_int16 = "int16"
+_int32 = "int32"
+_int64 = "int64"
+_bfloat16 = "bfloat16"
+_float16 = "float16"
+_float32 = "float32"
+_float64 = "float64"
+_complex64 = "complex64"
+_complex128 = "complex128"
+
+
+def _make_dtype(_dtypes, name):
+    class _Cls(AbstractDtype):
+        dtypes = _dtypes
+
+    _Cls.__name__ = name
+    _Cls.__qualname__ = name
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        _Cls.__module__ = "builtins"
+    else:
+        _Cls.__module__ = "jaxtyping"
+    return _Cls
+
+
+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]
+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
+
+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")
+Real = _make_dtype(floats + uints + ints, "Real")
+Num = _make_dtype(uints + ints + floats + complexes, "Num")
+
+Shaped = _make_dtype(_any_dtype, "Shaped")
+
+if has_jax:
+    Key = _make_dtype(_prng_key, "Key")
+    # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
+    # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype `uint32`.
+    PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
+    Scalar = Shaped[jax.Array, ""]
+    ScalarLike = Shaped[jax.typing.ArrayLike, ""]

jaxtyping/_decorator.py

@@ -0,0 +1,736 @@
+# 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 functools as ft
+import inspect
+import itertools as it
+import sys
+import types
+import weakref
+from typing import Any, get_args, get_origin, get_type_hints, overload
+
+
+try:
+    import jax._src.traceback_util as traceback_util
+except ImportError:
+    pass
+else:
+    traceback_util.register_exclusion(__file__)
+
+
+from ._storage import pop_shape_memo, push_shape_memo
+
+
+_jaxtyped_fns = weakref.WeakSet()
+
+
+class TypeCheckError(TypeError):
+    pass
+
+
+TypeCheckError.__module__ = "jaxtyping"  # appears in error messages
+
+
+@overload
+def jaxtyped(*, typechecker=None):
+    ...
+
+
+@overload
+def jaxtyped(fn, *, typechecker=None):
+    ...
+
+
+def jaxtyped(fn=None, *, typechecker=None):
+    """Decorate a function with this to perform runtime type-checking of its arguments
+    and return value. Decorate a dataclass to perform type-checking of its attributes.
+
+    !!! Example
+
+        ```python
+        # Import both the annotation and the `jaxtyped` decorator from `jaxtyping`
+        from jaxtyping import Array, Float, jaxtyped
+
+        # Use your favourite typechecker: usually one of the two lines below.
+        from typeguard import typechecked as typechecker
+        from beartype import beartype as typechecker
+
+        # Type-check a function
+        @jaxtyped(typechecker=typechecker)
+        def batch_outer_product(x: Float[Array, "b c1"],
+                                y: Float[Array, "b c2"]
+                              ) -> Float[Array, "b c1 c2"]:
+            return x[:, :, None] * y[:, None, :]
+
+        # Type-check a dataclass
+        @jaxtyped(typechecker=typechecker)
+        @dataclass
+        class MyDataclass:
+            x: int
+            y: Float[Array "b c"]
+        ```
+
+    **Arguments:**
+
+    - `fn`: The function or dataclass to decorate.
+    - `typechecker`: The runtime type-checker to use. This should be a function
+        decorator that will raise an exception if there is a type error, e.g.
+        ```python
+        @typechecker
+        def f(x: int):
+            pass
+
+        f("a string is not an integer")  # this line should raise an exception
+        ```
+        Common choices are `typechecker=beartype.beartype` or
+        `typechecker=typeguard.typechecked`. Can also be set as `typechecker=None` to
+        skip automatic runtime type-checking, but still support manual `isinstance`
+        checks inside the function body:
+        ```python
+        @jaxtyped
+        def f(x):
+            assert isinstance(x, Float[Array, "batch channel"])
+        ```
+
+    **Returns:**
+
+    If `fn` is a function (including a `staticmethod`, `classmethod`, or `property`),
+    then a wrapped function is returned.
+
+    If `fn` is a dataclass, then `fn` is returned directly, and additionally its
+    `__init__` method is wrapped and modified in-place.
+
+    !!! Info "Old syntax"
+
+        jaxtyping previously (before v0.2.24) recommended using this double-decorator
+        syntax:
+        ```python
+        @jaxtyped
+        @typechecker
+        def f(...): ...
+        ```
+        This is still supported, but the `jaxtyped(typechecker=typechecker)` syntax
+        discussed above will produce easier-to-debug error messages. Under the hood, the
+        new syntax more carefully manipulates the typechecker so as to determine where
+        a type-check error arises.
+
+    ??? Info "Notes for advanced users"
+
+        **Dynamic contexts:**
+
+        Put precisely, the axis names in e.g. `Float[Array, "batch channels"]` and the
+        structure names in e.g. `PyTree[int, "T"]` are all scoped to the thread-local
+        dynamic context of a `jaxtyped`-wrapped function. If from within that function
+        we then call another `jaxtyped`-wrapped function, then a new context is pushed
+        to the stack. The axis sizes and PyTree structures of this inner function will
+        then not be compared against the axis sizes and PyTree structures of the outer
+        function. After the inner function returns then this inner context is popped
+        from the stack, and the previous context is returned to.
+
+        **isinstance:**
+
+        Binding of a value against a name is done with an `isinstance` check, for
+        example `isinstance(jnp.zeros((3, 4)), Float[Array, "dim1 dim2"])` will bind
+        `dim1=3` and `dim2=4`. In practice these `isinstance` checks are usually done by
+        the run-time typechecker `typechecker` that is supplied as an argument.
+
+        This can also be done manually: add `isinstance` checks inside a function body
+        and they will contribute to the same collection of consistency checks as are
+        performed by the typechecker on the arguments and return values. (Or you can
+        forgo such a typechecker altogether -- i.e. `typechecker=None` -- and only do
+        your own manual `isinstance` checks.)
+
+        Only `isinstance` checks that pass will contribute to the store of values; those
+        that fail will not. As such it is safe to write e.g.
+        `assert not isinstance(x, Float32[Array, "foo"])`.
+
+        **Decoupling contexts from function calls:**
+
+        If you would like a new dynamic context *without* calling a new function, then
+        `jaxtyped` may be passed the string `"context"` and used as a context manager:
+        ```python
+        with jaxtyped("context"):
+            assert isinstance(x, Float[Array, "batch channel"])
+        ```
+        which is equivalent to placing this code inside a new function wrapped in
+        `jaxtyped(typechecker=None)`. Usage like this is very rare; it's mostly only
+        useful when working at the global scope.
+
+        Conversely, if you would like to call a new function *without* creating a new
+        dynamic context (and using the same set of axis and structure values), then
+        simply do not add a `jaxtyped` decorator to your inner function, whilst
+        continuing to perform type-checking in whatever way you prefer.
+    """
+
+    if fn is None:
+        return ft.partial(jaxtyped, typechecker=typechecker)
+    elif type(fn) is types.FunctionType and fn in _jaxtyped_fns:
+        return fn
+    elif inspect.isclass(fn):
+        if dataclasses.is_dataclass(fn) and typechecker is not None:
+            # This does not check that the arguments passed to `__init__` match the
+            # type annotations. There may be a custom user `__init__`, or a
+            # dataclass-generated `__init__` used alongside
+            # `equinox.field(converter=...)`
+
+            init = fn.__init__
+
+            @ft.wraps(init)
+            def __init__(self, *args, **kwargs):
+                init(self, *args, **kwargs)
+                # `fn.__init__` is late-binding to the `__init__` function that
+                # we're in now. (Or to someone else's monkey-patch.) Either way,
+                # this checks that we're in the "top-level" `__init__`, and not one
+                # that is being called via `super()`. We don't want to trigger too
+                # early, before all fields have been assigned.
+                #
+                # We're not checking `if self.__class__ is fn` because Equinox
+                # replaces the with a defrozen version of itself during `__init__`,
+                # so the check wouldn't trigger.
+                #
+                # We're not doing this check by adding it to the end of the
+                # metaclass `__call__`, because Python doesn't allow you
+                # monkey-patch metaclasses.
+                if self.__class__.__init__ is fn.__init__:
+                    _check_dataclass_annotations(self, typechecker)
+
+            fn.__init__ = __init__
+        return fn
+    # 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__, typechecker=typechecker))
+    elif isinstance(fn, staticmethod):
+        return staticmethod(jaxtyped(fn.__func__, typechecker=typechecker))
+    elif isinstance(fn, property):
+        if fn.fget is None:
+            fget = None
+        else:
+            fget = jaxtyped(fn.fget, typechecker=typechecker)
+        if fn.fset is None:
+            fset = None
+        else:
+            fset = jaxtyped(fn.fset, typechecker=typechecker)
+        if fn.fdel is None:
+            fdel = None
+        else:
+            fdel = jaxtyped(fn.fdel, typechecker=typechecker)
+        return property(fget=fget, fset=fset, fdel=fdel)
+    elif fn == "context":
+        if typechecker is not None:
+            raise ValueError(
+                "Cannot use `jaxtyped` as a context with a typechecker. That is, "
+                "`with jaxtyped('context', typechecker=...):`. is not allowed. In this "
+                "case the type checker does not actually do anything, as there is no "
+                "function to type-check."
+            )
+        return _JaxtypingContext()
+    else:
+        if typechecker is None:
+            # Probably being used in the old style as
+            # ```
+            # @jaxtyped
+            # @typechecker
+            # def foo(x: int): ...
+            # ```
+            # in which case make a best-effort attempt to add shape information for any
+            # type errors.
+
+            signature = inspect.signature(fn)
+
+            @ft.wraps(fn)
+            def wrapped_fn(*args, **kwargs):  # pyright: ignore
+                bound = signature.bind(*args, **kwargs)
+                memos = push_shape_memo(bound.arguments)
+                try:
+                    return fn(*args, **kwargs)
+                except Exception as e:
+                    if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
+                        shape_info = _exc_shape_info(memos)
+                        if shape_info != "":
+                            msg = (
+                                "The preceding error occurred within the scope of a "
+                                "`jaxtyping.jaxtyped` function, and may be due to a "
+                                "typecheck error. "
+                            )
+                            e.add_note(_jaxtyping_note_str(_spacer + msg + shape_info))
+                    raise
+                finally:
+                    pop_shape_memo()
+
+        else:
+            # New-style
+            # ```
+            # @jaxtyped(typechecker=typechecker)
+            # def foo(x: int): ...
+            # ```
+            # in which case we can do a better job reporting errors.
+
+            full_signature = inspect.signature(fn)
+            try:
+                destring_annotations = get_type_hints(fn, include_extras=True)
+            except NameError:
+                # Best-effort attempt to destringify annotations.
+                pass
+            else:
+                new_params = []
+                for p_name, p_value in full_signature.parameters.items():
+                    p_annotation = destring_annotations.get(p_name, p_value.annotation)
+                    p_value = p_value.replace(annotation=p_annotation)
+                    new_params.append(p_value)
+                return_annotation = destring_annotations.get(
+                    "return", full_signature.return_annotation
+                )
+                full_signature = full_signature.replace(
+                    parameters=new_params, return_annotation=return_annotation
+                )
+
+            param_signature = full_signature.replace(
+                return_annotation=inspect.Signature.empty
+            )
+            module = getattr(fn, "__module__", "generated")
+
+            full_fn, output_name = _make_fn_with_signature(
+                "check_return", full_signature, module, output=True
+            )
+            full_fn = typechecker(full_fn)
+
+            param_fn = _make_fn_with_signature(
+                "check_params", param_signature, module, output=False
+            )
+            param_fn = typechecker(param_fn)
+
+            @ft.wraps(fn)
+            def wrapped_fn(*args, **kwargs):
+                # Raise bind-time errors before we do any shape analysis. (I.e. skip
+                # the pointless jaxtyping information for a non-typechecking failure.)
+                bound = param_signature.bind(*args, **kwargs)
+
+                memos = push_shape_memo(bound.arguments)
+                try:
+                    # First type-check just the parameters before the function is
+                    # called.
+                    try:
+                        param_fn(*args, **kwargs)
+                    except Exception as e:
+                        if hasattr(e, "_jaxtyping_malformed"):
+                            raise
+                        else:
+                            argmsg = _get_problem_arg(
+                                param_signature, args, kwargs, module, typechecker
+                            )
+                            try:
+                                name = fn.__name__
+                            except AttributeError:
+                                name = fn.__class__.__name__
+                            param_values = _pformat(bound.arguments, short_self=True)
+                            param_hints = _remove_typing(param_signature)
+                            msg = (
+                                "Type-check error whilst checking the parameters of "
+                                f"{name}.{argmsg}\n"
+                                f"Called with arguments: {param_values}\n"
+                                f"Parameter annotations: {param_hints}.\n"
+                                + _exc_shape_info(memos)
+                            )
+                            raise TypeCheckError(msg) from e
+
+                    # Actually call the function.
+                    out = fn(*args, **kwargs)
+
+                    if full_signature.return_annotation is not inspect.Signature.empty:
+                        # Now type-check the return value. We need to include the
+                        # parameters in the type-checking here in case there are any
+                        # type variables shared across the parameters and return.
+                        #
+                        # Incidentally this does mean that if `fn` mutates its arguments
+                        # so that they no longer satisfy their type annotations, this
+                        # will throw an error here. But that's like, super weird, so
+                        # don't do that. An error in that scenario is probably still
+                        # desirable.
+                        #
+                        # There is a small performance concern here when used in
+                        # non-jit'd contexts, like PyTorch, due to the duplicate
+                        # checking of the parameters. Unfortunately there doesn't seem
+                        # to be a way around that, so c'est la vie.
+                        kwargs[output_name] = out
+                        try:
+                            full_fn(*args, **kwargs)
+                        except Exception as e:
+                            if hasattr(e, "_jaxtyping_malformed"):
+                                raise
+                            else:
+                                try:
+                                    name = fn.__name__
+                                except AttributeError:
+                                    name = fn.__class__.__name__
+                                param_values = _pformat(
+                                    bound.arguments, short_self=True
+                                )
+                                return_value = _pformat(out, short_self=False)
+                                param_hints = _remove_typing(param_signature)
+                                return_hint = _remove_typing(
+                                    full_signature.return_annotation
+                                )
+                                if return_hint.startswith(
+                                    "<class '"
+                                ) and return_hint.endswith("'>"):
+                                    return_hint = return_hint[8:-2]
+                                msg = (
+                                    "Type-check error whilst checking the return value "
+                                    f"of {name}.\n"
+                                    f"Called with arguments: {param_values}\n"
+                                    f"Return value: {return_value}\n"
+                                    f"Parameter annotations: {param_hints}.\n"
+                                    f"Return annotation: {return_hint}.\n"
+                                    + _exc_shape_info(memos)
+                                )
+                                raise TypeCheckError(msg) from e
+
+                    return out
+                finally:
+                    pop_shape_memo()
+
+        _jaxtyped_fns.add(wrapped_fn)
+        return wrapped_fn
+
+
+class _JaxtypingContext:
+    def __enter__(self):
+        push_shape_memo({})
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        pop_shape_memo()
+
+
+def _check_dataclass_annotations(self, typechecker):
+    """Creates and calls a function that checks the attributes of `self`
+
+    `self` should be a dataclass instancae. `typechecker` should be e.g.
+    `beartype.beartype` or `typeguard.typechecked`.
+    """
+    parameters = [inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)]
+    values = {}
+    for field in dataclasses.fields(self):
+        annotation = field.type
+        if isinstance(annotation, str):
+            # Don't check stringified annotations. These are basically impossible to
+            # resolve correctly, so just skip them.
+            continue
+        if get_origin(annotation) is type:
+            args = get_args(annotation)
+            if len(args) == 1 and isinstance(args[0], str):
+                # We also special-case this one kind of partially-stringified type
+                # annotation, so as to support Equinox <v0.11.1.
+                # This was fixed in Equinox in
+                # https://github.com/patrick-kidger/equinox/pull/543
+                continue
+        try:
+            value = getattr(self, field.name)  # noqa: F841
+        except AttributeError:
+            continue  # allow uninitialised fields, which are allowed on dataclasses
+
+        parameters.append(
+            inspect.Parameter(
+                field.name,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=field.type,
+            )
+        )
+        values[field.name] = value
+
+    signature = inspect.Signature(parameters)
+    module = self.__class__.__module__
+    f = _make_fn_with_signature(
+        self.__class__.__name__, signature, module, output=False
+    )
+    f = jaxtyped(f, typechecker=typechecker)
+    f(self, **values)
+
+
+def _make_fn_with_signature(
+    name: str, signature: inspect.Signature, module: str, output: bool
+):
+    """Dynamically creates a function `fn` with name `name` and signature `signature`.
+
+    If `output=True` then `fn` will consume an additional keyword-only argument (in
+    addition to the provided signature), and will directly return this argument. In this
+    case the returned value from `_make_fn_with_signature` is a 2-tuple of `(fn, name)`,
+    where `fn` is the generated function, and `name` is the name of this extra argument.
+
+    If `output=False` then `fn` will just have a single `pass` statement, and the
+    returned value from `_make_fn_with_signature` will just be `fn`.
+
+    ---
+
+    Note that this function operates by dynamically creating and eval'ing a string, not
+    simply by assigning `__signature__` and `__annotations__`. The latter is enough for
+    typeguard (at least v2), but does not work with beartype (at least v16).
+    """
+    pos = []
+    pos_or_key = []
+    varpos = []
+    key = []
+    varkey = []
+    for p in signature.parameters.values():
+        if p.kind == inspect.Parameter.POSITIONAL_ONLY:
+            pos.append(p)
+        elif p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
+            pos_or_key.append(p)
+        elif p.kind == inspect.Parameter.VAR_POSITIONAL:
+            varpos.append(p)
+        elif p.kind == inspect.Parameter.KEYWORD_ONLY:
+            key.append(p)
+        elif p.kind == inspect.Parameter.VAR_KEYWORD:
+            varkey.append(p)
+        else:
+            assert False
+
+    param_names = frozenset(signature.parameters.keys())
+    if output:
+        output_name = _gensym(param_names, prefix="ret")
+        outstr = "return " + output_name
+        param_names = param_names | frozenset({output_name})
+        key.append(inspect.Parameter(output_name, kind=inspect.Parameter.KEYWORD_ONLY))
+    else:
+        outstr = "pass"
+
+    scope = {name: None}
+    name_to_annotation = {}
+    name_to_default = {}
+    param_triples = (
+        (p.name, p.annotation, p.default) for p in signature.parameters.values()
+    )
+    if output:
+        triples = it.chain(
+            param_triples,
+            [
+                ("return", signature.return_annotation, inspect.Signature.empty),
+                (output_name, Any, inspect.Signature.empty),
+            ],
+        )
+    else:
+        triples = it.chain(
+            param_triples,
+            [("return", signature.return_annotation, inspect.Signature.empty)],
+        )
+    for p_name, p_annotation, p_default in triples:
+        annotation_name = _gensym(frozenset(scope.keys()) | param_names, prefix="T")
+        name_to_annotation[p_name] = annotation_name
+        if p_annotation is inspect.Signature.empty or isinstance(p_annotation, str):
+            # If we have a stringified annotation here it's because the get_type_hints
+            # lookup above failed. Typically this occurs when using a local variable as
+            # the annotation. In this case we really have no idea what the annotation
+            # refers to, so just set it to Any.
+            # This does mean that we don't handle partially-stringified local
+            # annotations, e.g. `type["Foo"]` for some local type `Foo`. Those will
+            # probably just error out. Nothing better we can do about that
+            # unfortunately.
+            scope[annotation_name] = Any
+        else:
+            scope[annotation_name] = p_annotation
+        default_name = _gensym(frozenset(scope.keys()) | param_names, prefix="default")
+        name_to_default[p_name] = default_name
+        scope[default_name] = p_default
+
+    argstr_pieces = []
+    if len(pos) > 0:
+        for p in pos:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+        argstr_pieces.append("/")
+    if len(pos_or_key) > 0:
+        for p in pos_or_key:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+    if len(varpos) == 1:
+        [p] = varpos
+        argstr_pieces.append(
+            "*" + _make_argpiece(p, name_to_annotation, name_to_default)
+        )
+    else:
+        assert len(varpos) == 0
+        if len(key) > 0:
+            argstr_pieces.append("*")
+    if len(key) > 0:
+        for p in key:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+    if len(varkey) == 1:
+        [p] = varkey
+        argstr_pieces.append(
+            "**" + _make_argpiece(p, name_to_annotation, name_to_default)
+        )
+    else:
+        assert len(varkey) == 0
+    argstr = ", ".join(argstr_pieces)
+
+    if signature.return_annotation is inspect.Signature.empty:
+        retstr = ""
+    else:
+        retstr = f"-> {name_to_annotation['return']}"
+
+    fnstr = f"def {name}({argstr}){retstr}:\n    {outstr}"
+    exec(fnstr, scope)
+    fn = scope[name]
+    fn.__module__ = module
+    assert fn is not None
+    if output:
+        return fn, output_name
+    else:
+        return fn
+
+
+def _gensym(names: frozenset[str], prefix: str) -> str:
+    assert prefix.isidentifier()
+    output_index = 0
+    output_name = prefix + str(output_index)
+    while output_name in names:
+        output_index += 1
+        output_name = prefix + str(output_index)
+    assert output_name.isidentifier()
+    return output_name
+
+
+def _make_argpiece(p, name_to_annotation, name_to_default):
+    if p.default is inspect.Signature.empty:
+        return f"{p.name}: {name_to_annotation[p.name]}"
+    else:
+        return f"{p.name}: {name_to_annotation[p.name]} = {name_to_default[p.name]}"
+
+
+def _get_problem_arg(
+    param_signature: inspect.Signature, args, kwargs, module, typechecker
+) -> str:
+    """Determines which argument was likely to be the problematic one responsible for
+    raising a type-check error.
+    """
+    # No performance concerns, as this is only used when we're about to raise an error
+    # anyway.
+    for keep_name in param_signature.parameters.keys():
+        new_parameters = []
+        for p_name, p in param_signature.parameters.items():
+            if p_name == keep_name:
+                new_parameters.append(
+                    inspect.Parameter(p.name, p.kind, annotation=p.annotation)
+                )
+            else:
+                new_parameters.append(inspect.Parameter(p.name, p.kind))
+        new_signature = inspect.Signature(new_parameters)
+        fn = _make_fn_with_signature(
+            "check_single_arg", new_signature, module, output=False
+        )
+        fn = typechecker(fn)  # but no `jaxtyped`; keep the same environment.
+        try:
+            fn(*args, **kwargs)
+        except Exception:
+            return f"\nThe problem arose whilst typechecking argument '{keep_name}'."
+    else:
+        # Could not localise the problem to a single argument -- probably due to
+        # e.g. a mismatched typevar, which each individual argument is okay with.
+        return ""
+
+
+def _remove_typing(x):
+    x = str(x)
+    x = x.replace(" jaxtyping.", " ")
+    x = x.replace("[jaxtyping.", "[")
+    x = x.replace("'jaxtyping.", "'")
+    x = x.replace(" typing.", " ")
+    x = x.replace("[typing.", "[")
+    x = x.replace("'typing.", "'")
+    return x
+
+
+def _pformat(x, short_self: bool):
+    # No performance concerns from delayed imports -- this is only used when we're about
+    # to raise an error anyway.
+    try:
+        # TODO(kidger): this is pretty ugly. We have a circular dependency
+        # equinox->jaxtyping->equinox. We could consider moving all the pretty-printing
+        # code from equinox into jaxtyping maybe? Or into some shared dependency?
+        import equinox as eqx
+
+        pformat = eqx.tree_pformat
+        if short_self:
+            try:
+                self = x["self"]
+            except KeyError:
+                pass
+            else:
+                is_self = lambda y: y is self
+                pformat = ft.partial(pformat, truncate_leaf=is_self)
+    except Exception:
+        import pprint
+
+        pformat = ft.partial(pprint.pformat, indent=2, compact=True)
+    return pformat(x)
+
+
+def _exc_shape_info(memos) -> str:
+    """Gives debug information on the current state of jaxtyping's internal memos.
+    Used in type-checking error messages.
+    """
+    single_memo, variadic_memo, pytree_memo, _ = memos
+    single_memo = {
+        name: size
+        for name, size in single_memo.items()
+        if not name.startswith("~~delete~~")
+    }
+    variadic_memo = {
+        name: shape
+        for name, (_, shape) in variadic_memo.items()
+        if not name.startswith("~~delete~~")
+    }
+    pieces = []
+    if len(single_memo) > 0 or len(variadic_memo) > 0:
+        pieces.append(
+            "The current values for each jaxtyping axis annotation are as follows."
+        )
+        for name, size in single_memo.items():
+            pieces.append(f"{name}={size}")
+        for name, shape in variadic_memo.items():
+            pieces.append(f"{name}={shape}")
+    if len(pytree_memo) > 0:
+        pieces.append(
+            "The current values for each jaxtyping PyTree structure annotation are as "
+            "follows."
+        )
+        for name, structure in pytree_memo.items():
+            pieces.append(f"{name}={structure}")
+    return "\n".join(pieces)
+
+
+class _jaxtyping_note_str(str):
+    """Used with `_no_jaxtyping_note` to flag that a note came from jaxtyping."""
+
+
+def _no_jaxtyping_note(e: Exception) -> bool:
+    """Checks if any of the exception's notes are from jaxtyping."""
+    try:
+        notes = e.__notes__
+    except AttributeError:
+        return True
+    else:
+        for note in notes:
+            if isinstance(note, _jaxtyping_note_str):
+                return False
+        return True
+
+
+_spacer = "--------------------\n"

jaxtyping/_import_hook.py

@@ -0,0 +1,427 @@
+# Copyright (c) 2022 Google LLC
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to
+# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+# the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+# This source code is adapted from typeguard:
+# https://github.com/agronholm/typeguard/blob/0dd7f7510b7c694e66a0d17d1d58d185125bad5d/src/typeguard/importhook.py
+#
+# Copied and adapted in compliance with the terms of typeguard's MIT license.
+# The original license is reproduced here.
+#
+# ---------
+#
+# This is the MIT license: http://www.opensource.org/licenses/mit-license.php
+#
+# Copyright (c) Alex Grönholm
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this
+# software and associated documentation files (the "Software"), to deal in the Software
+# without restriction, including without limitation the rights to use, copy, modify,
+# merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
+# permit persons to whom the Software is furnished to do so, subject to the following
+# conditions:
+#
+# The above copyright notice and this permission notice shall be included in all copies
+# or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+# PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
+# CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE
+# OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+#
+# ---------
+
+
+import ast
+import functools as ft
+import hashlib
+import sys
+from collections.abc import Sequence
+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 Optional, Union
+from unittest.mock import patch
+
+
+# The name of this function is magical
+def _call_with_frames_removed(f, *args, **kwargs):
+    return f(*args, **kwargs)
+
+
+def _optimized_cache_from_source(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.
+    # Version 7: Using the same md5 hash of the `typechecker` argument
+    #    for importlib and decorator lookup.
+    # Version 8: Now using new-style `jaxtyped(typechecker=...)` rather than old-style
+    #    double-decorators.
+    return cache_from_source(
+        path, debug_override, optimization=f"jaxtyping8{typechecker_hash}"
+    )
+
+
+class Typechecker:
+    lookup = {}
+
+    def __init__(self, typechecker):
+        self.ast = None
+
+        if isinstance(typechecker, str):
+            # If the typechecker is a string, then we parse it
+            string_to_eval = (
+                "def f(x, *args, **kwargs):\n"
+                + f"  import {typechecker.split('.', 1)[0]}\n"
+                + f"  return {typechecker}(x, *args, **kwargs)"
+            )
+
+            # md5 hashing instead of __hash__
+            # because __hash__ is different for each Python session
+            self.hash = hashlib.md5(typechecker.encode("utf-8")).hexdigest()
+
+            vars = {}
+            exec(string_to_eval, {}, vars)
+            Typechecker.lookup[self.hash] = vars["f"]
+
+        elif typechecker is None:
+            # If it is None, ignore it silently (use dummy decorator)
+            self.hash = "0"
+            Typechecker.lookup[self.hash] = lambda x, *_, **__: x
+        else:
+            # Passed typechecker is invalid
+            raise TypeError(
+                "Jaxtyping typechecker has to be either a string or a None."
+            )
+
+    def get_hash(self):
+        return self.hash
+
+    def get_ast(self):
+        # we compile AST only if we missed importlib cache
+        if self.ast is None:
+            self.ast = (
+                ast.parse(
+                    f"@jaxtyping.jaxtyped(typechecker=jaxtyping._import_hook.Typechecker.lookup['{self.hash}'])\n"
+                    "def _():\n    ..."
+                )
+                .body[0]
+                .decorator_list[0]
+            )
+
+        return self.ast
+
+
+class JaxtypingTransformer(ast.NodeVisitor):
+    def __init__(self, *, typechecker: Typechecker) -> None:
+        self._parents: list[ast.AST] = []
+        self._typechecker = typechecker
+
+    def visit_Module(self, node: ast.Module):
+        # Insert "import jaxtyping" after any "from __future__ ..." imports
+        for i, child in enumerate(node.body):
+            if isinstance(child, ast.ImportFrom) and child.module == "__future__":
+                continue
+            elif isinstance(child, ast.Expr) and isinstance(child.value, ast.Str):
+                continue  # module docstring
+            else:
+                node.body.insert(i, ast.Import(names=[ast.alias("jaxtyping", None)]))
+                break
+
+        self._parents.append(node)
+        self.generic_visit(node)
+        self._parents.pop()
+        return node
+
+    def visit_ClassDef(self, node: ast.ClassDef):
+        # Place at the start of the decorator list, so that `@dataclass` decorators get
+        # called first.
+        node.decorator_list.insert(0, self._typechecker.get_ast())
+        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 end of the decorator list, because:
+            # - as otherwise we wrap e.g. `jax.custom_{jvp,vjp}` and lose the ability
+            #     to `defjvp` etc.
+            # - decorators frequently remove annotations from functions, and we'd like
+            #     to use those annotations.
+            # - typeguard in particular wants to be at the end of the decorator list, as
+            #     it works by recompling the wrapped function.
+            #
+            # 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.
+            node.decorator_list.append(self._typechecker.get_ast())
+
+        self._parents.append(node)
+        self.generic_visit(node)
+        self._parents.pop()
+        return node
+
+
+class _JaxtypingLoader(SourceFileLoader):
+    def __init__(self, *args, typechecker: Typechecker, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._typechecker = typechecker
+
+    def source_to_code(self, data, path, *, _optimize=-1):
+        source = decode_source(data)
+        tree = _call_with_frames_removed(
+            compile,
+            source,
+            path,
+            "exec",
+            ast.PyCF_ONLY_AST,
+            dont_inherit=True,
+            optimize=_optimize,
+        )
+        tree = JaxtypingTransformer(typechecker=self._typechecker).visit(tree)
+        ast.fix_missing_locations(tree)
+        return _call_with_frames_removed(
+            compile, tree, path, "exec", dont_inherit=True, optimize=_optimize
+        )
+
+    def exec_module(self, module):
+        # Use a custom optimization marker - the import lock should make this monkey
+        # patch safe
+        with patch(
+            "importlib._bootstrap_external.cache_from_source",
+            ft.partial(_optimized_cache_from_source, self._typechecker.get_hash()),
+        ):
+            return super().exec_module(module)
+
+
+class _JaxtypingFinder(MetaPathFinder):
+    """Wraps another path finder and instruments the module with `@jaxtyped` and
+    `@typechecked` if `should_instrument()` returns `True`.
+
+    Should not be used directly, but rather via `install_import_hook`.
+    """
+
+    def __init__(self, modules, original_pathfinder, typechecker: Typechecker):
+        self.modules = modules
+        self._original_pathfinder = original_pathfinder
+        self._typechecker = typechecker
+
+    def find_spec(self, fullname, path=None, target=None):
+        if self.should_instrument(fullname):
+            spec = self._original_pathfinder.find_spec(fullname, path, target)
+            if spec is not None and isinstance(spec.loader, SourceFileLoader):
+                spec.loader = _JaxtypingLoader(
+                    spec.loader.name, spec.loader.path, typechecker=self._typechecker
+                )
+                return spec
+
+        return None
+
+    def should_instrument(self, module_name: str) -> bool:
+        """Determine whether the module with the given name should be instrumented.
+
+        **Arguments:**
+
+        - `module_name`: the full name of the module that is about to be imported
+            (e.g. ``xyz.abc``)
+        """
+        for module in self.modules:
+            if module_name == module or module_name.startswith(module + "."):
+                return True
+
+        return False
+
+
+class ImportHookManager:
+    def __init__(self, hook: MetaPathFinder):
+        self.hook = hook
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.uninstall()
+
+    def uninstall(self):
+        try:
+            sys.meta_path.remove(self.hook)
+        except ValueError:
+            pass  # already removed
+
+
+# Deliberately no default for `typechecker` so that folks must opt-in to not having
+# a typechecker.
+def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optional[str]):
+    """Automatically apply the `@jaxtyped(typechecker=typechecker)` decorator to every
+    function and dataclass over a whole codebase.
+
+    !!! Tip "Usage"
+
+        ```python
+        from jaxtyping import install_import_hook
+        # Plus any one of the following:
+
+        # decorate `@jaxtyped(typechecker=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 and
+                                # dataclasses.
+
+        # decorate `@jaxtyped(typechecker=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"], ...)
+        ```
+
+    **Arguments:**:
+
+    - `modules`: the names of the modules in which to automatically apply `@jaxtyped`.
+    - `typechecker`: the module and function of the typechecker you want to use, as a
+        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: end-user script"
+
+        ```python
+        ### entry_point.py
+        from jaxtyping import install_import_hook
+        with install_import_hook("main", "typeguard.typechecked"):
+            import main
+
+        ### 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.
+        ```
+
+    !!! warning
+
+        Stringified dataclass annotations, e.g.
+        ```python
+        @dataclass()
+        class Foo:
+            x: "int"
+        ```
+        will be silently skipped without checking them. This is because these are
+        essentially impossible to resolve at runtime. Such stringified annotations
+        typically occur either when using them for forward references, or when using
+        `from __future__ import annotations`. (You should never use the latter, it is
+        largely incompatible with runtime type checking.)
+
+        Partially stringified dataclass annotations, e.g.
+        ```python
+        @dataclass()
+        class Foo:
+            x: tuple["int"]
+        ```
+        will likely raise an error, and must not be used at all.
+
+    !!! warning
+
+        If a function already has any decorators on it, then `@jaxtyped` will get added
+        at the bottom of the decorator list, e.g.
+        ```python
+        @some_other_decorator
+        @jaxtyped(typechecker=beartype.beartype)
+        def foo(...): ...
+        ```
+        This is to support the common case in which
+        `some_other_decorator = jax.custom_jvp` etc.
+
+        If a class already has any decorators in it, then `@jaxtyped` will get added to
+        the top of the decorator list, e.g.
+        ```python
+        @jaxtyped(typechecker=beartype.beartype)
+        @some_other_decorator
+        class A:
+            ...
+        ```
+        This is to support the common case in which
+        `some_other_decorator = dataclasses.dataclass`.
+    """  # noqa: E501
+
+    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)
+            and finder.__name__ == "PathFinder"
+            and hasattr(finder, "find_spec")
+        ):
+            break
+    else:
+        raise RuntimeError("Cannot find a PathFinder in sys.meta_path")
+
+    wrapped_typechecker = Typechecker(typechecker)
+    hook = _JaxtypingFinder(modules, finder, wrapped_typechecker)
+    sys.meta_path.insert(0, hook)
+    return ImportHookManager(hook)

jaxtyping/_indirection.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2022 Google LLC
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to
+# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+# the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+# Note that `from typing 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 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 Real,  # 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 PRNGKeyArray,  # noqa: F401
+    Array as Scalar,  # noqa: F401
+)
+from jax.typing import ArrayLike as ScalarLike  # noqa: F401

jaxtyping/_ipython_extension.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2022 Google LLC
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to
+# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+# the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+from ._import_hook import JaxtypingTransformer, Typechecker
+
+
+try:
+    from IPython.core.magic import line_magic, Magics, magics_class
+
+    @magics_class
+    class ChooseTypecheckerMagics(Magics):
+        @line_magic("jaxtyping.typechecker")
+        def typechecker(self, typechecker):
+            # remove old JaxtypingTransformer, if present
+            self.shell.ast_transformers = list(
+                filter(
+                    lambda x: not isinstance(x, JaxtypingTransformer),
+                    self.shell.ast_transformers,
+                )
+            )
+
+            # add new one
+            self.shell.ast_transformers.append(
+                JaxtypingTransformer(typechecker=Typechecker(typechecker))
+            )
+
+except Exception:
+    # Very broad exception-handling, as e.g. IPython will sometimes be
+    # present but fail to import for mysterious reasons.
+    pass
+
+
+def load_ipython_extension(ipython):
+    try:
+        ipython.register_magics(ChooseTypecheckerMagics)
+    except NameError:
+        raise NameError(
+            "ChooseTypecheckerMagics is not defined.\n\n"
+            + "You may be trying to use IPython extension without IPython installed."
+        )

jaxtyping/_pytest_plugin.py

@@ -19,7 +19,7 @@
 
 import sys
 
-from .import_hook import install_import_hook
+from ._import_hook import install_import_hook
 
 
 def pytest_addoption(parser):
@@ -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

@@ -0,0 +1,325 @@
+# Copyright (c) 2022 Google LLC
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to
+# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+# the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+import functools as ft
+import typing
+from typing import Any, Generic, TypeVar
+
+import jax.tree_util as jtu
+import typeguard
+
+from ._raise import jaxtyping_raise_from
+from ._storage import (
+    clear_treeflatten_memo,
+    clear_treepath_memo,
+    get_shape_memo,
+    set_shape_memo,
+    set_treeflatten_memo,
+    set_treepath_memo,
+)
+
+
+_T = TypeVar("_T")
+
+
+class _FakePyTree(Generic[_T]):
+    pass
+
+
+_FakePyTree.__name__ = "PyTree"
+_FakePyTree.__qualname__ = "PyTree"
+_FakePyTree.__module__ = "builtins"
+
+
+class _MetaPyTree(type):
+    def __call__(self, *args, **kwargs):
+        raise RuntimeError("PyTree cannot be instantiated")
+
+    def __instancecheck__(cls, obj):
+        if not hasattr(cls, "leaftype"):
+            return True  # Just `isinstance(x, PyTree)`
+        # Handle beartype doing `isinstance(None, hint)` to check if
+        # is `instance`able.
+        if obj is None:
+            return True
+
+        single_memo, variadic_memo, pytree_memo, arg_memo = get_shape_memo()
+        single_memo_bak = single_memo.copy()
+        variadic_memo_bak = variadic_memo.copy()
+        pytree_memo_bak = pytree_memo.copy()
+        arg_memo_bak = arg_memo.copy()
+        try:
+            out = cls._check(obj, pytree_memo)
+        except Exception:
+            set_shape_memo(
+                single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
+            )
+            raise
+        if out:
+            return True
+        else:
+            set_shape_memo(
+                single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
+            )
+            return False
+
+    def _check(cls, obj, pytree_memo):
+        if cls.leaftype is Any:
+
+            def is_flatten_leaftype(x):
+                return False
+
+            def is_check_leaftype(x):
+                return True
+
+        else:
+            # We could use `isinstance` here but that would fail for more complicated
+            # types, e.g. PyTree[tuple[int]]. So at least internally we make a
+            # particular choice of typechecker.
+            #
+            # Deliberately not using @jaxtyped so that we share the same `memo` as
+            # whatever dynamic context we're currently in.
+            @typeguard.typechecked
+            def accepts_leaftype(x: cls.leaftype):
+                pass
+
+            def is_leaftype(x):
+                try:
+                    accepts_leaftype(x)
+                except _TypeCheckError:
+                    return False
+                else:
+                    return True
+
+            is_flatten_leaftype = is_check_leaftype = is_leaftype
+
+        set_treeflatten_memo()
+        try:
+            leaves, structure = jtu.tree_flatten(obj, is_leaf=is_flatten_leaftype)
+        finally:
+            clear_treeflatten_memo()
+        if cls.structure is not None:
+            if cls.structure.isidentifier():
+                try:
+                    prev_structure = pytree_memo[cls.structure]
+                except KeyError:
+                    pytree_memo[cls.structure] = structure
+                else:
+                    if prev_structure != structure:
+                        return False
+            else:
+                named_pytree = 0
+                pieces = cls.structure.split()
+                if pieces[0] == "...":
+                    pieces = pieces[1:]
+                    prefix = False
+                    suffix = True
+                elif pieces[-1] == "...":
+                    pieces = pieces[:-1]
+                    prefix = True
+                    suffix = False
+                else:
+                    prefix = False
+                    suffix = False
+                for identifier in pieces:
+                    try:
+                        prev_structure = pytree_memo[identifier]
+                    except KeyError as e:
+                        jaxtyping_raise_from(
+                            NameError(
+                                f"Cannot process composite structure '{cls.structure}' "
+                                f"as the structure name {identifier} has not been seen "
+                                "before."
+                            ),
+                            e,
+                        )
+                    # Not using `PyTreeDef.compose` due to JAX bug #18218.
+                    prev_pytree = jtu.tree_unflatten(
+                        prev_structure, [0] * prev_structure.num_leaves
+                    )
+                    named_pytree = jtu.tree_map(lambda _: prev_pytree, named_pytree)
+                named_structure = jtu.tree_structure(named_pytree)
+                if prefix:
+                    dummy_pytree = jtu.tree_unflatten(structure, [0] * len(leaves))
+                    dummy_named = jtu.tree_unflatten(
+                        named_structure, [0] * named_structure.num_leaves
+                    )
+                    try:
+                        jtu.tree_map(lambda _, __: 0, dummy_named, dummy_pytree)
+                    except ValueError:
+                        return False
+                elif suffix:
+                    has_structure = lambda x: jtu.tree_structure(x) == named_structure
+                    dummy_pytree = jtu.tree_unflatten(structure, [0] * len(leaves))
+                    dummy_leaves = jtu.tree_leaves(dummy_pytree, is_leaf=has_structure)
+                    if any(not has_structure(x) for x in dummy_leaves):
+                        return False
+                else:
+                    if structure != named_structure:
+                        return False
+
+        try:
+            for leaf_index, leaf in enumerate(leaves):
+                if cls.structure is not None:
+                    set_treepath_memo(leaf_index, cls.structure)
+                if not is_check_leaftype(leaf):
+                    return False
+                clear_treepath_memo()
+        finally:
+            clear_treepath_memo()
+        return True
+
+    # Can't return a generic (e.g. _FakePyTree[item]) because generic aliases don't do
+    # the custom __instancecheck__ that we want.
+    # We can't add that __instancecheck__  via subclassing, e.g.
+    # type("PyTree", (Generic[_T],), {}), because dynamic subclassing of typeforms
+    # isn't allowed.
+    # Likewise we can't do types.new_class("PyTree", (Generic[_T],), {}) because that
+    # has __module__ "types", e.g. we get types.PyTree[int].
+    @ft.lru_cache(maxsize=None)
+    def __getitem__(cls, item):
+        if isinstance(item, tuple):
+            if len(item) == 2:
+
+                class X(PyTree):
+                    leaftype = item[0]
+                    structure = item[1].strip()
+
+                if not isinstance(X.structure, str):
+                    raise ValueError(
+                        "The structure annotation `struct` in "
+                        "`jaxtyping.PyTree[leaftype, struct]` must be be a string, "
+                        f"e.g. `jaxtyping.PyTree[leaftype, 'T']`. Got '{X.structure}'."
+                    )
+                pieces = X.structure.split()
+                if len(pieces) == 0:
+                    raise ValueError(
+                        "The string `struct` in `jaxtyping.PyTree[leaftype, struct]` "
+                        "cannot be the empty string."
+                    )
+                for piece_index, piece in enumerate(pieces):
+                    if (piece_index == 0) or (piece_index == len(pieces) - 1):
+                        if piece == "...":
+                            continue
+                    if not piece.isidentifier():
+                        raise ValueError(
+                            "The string `struct` in "
+                            "`jaxtyping.PyTree[leaftype, struct]` must be be a "
+                            "whitespace-separated sequence of identifiers, e.g. "
+                            "`jaxtyping.PyTree[leaftype, 'T']` or "
+                            "`jaxtyping.PyTree[leaftype, 'foo bar']`.\n"
+                            "(Here, 'identifier' is used in the same sense as in "
+                            "regular Python, i.e. a valid variable name.)\n"
+                            f"Got piece '{piece}' in overall structure '{X.structure}'."
+                        )
+                name = str(_FakePyTree[item[0]])[:-1] + ', "' + item[1].strip() + '"]'
+            else:
+                raise ValueError(
+                    "The subscript `foo` in `jaxtyping.PyTree[foo]` must either be a "
+                    "leaf type, e.g. `PyTree[int]`, or a 2-tuple of leaf and "
+                    "structure, e.g. `PyTree[int, 'T']`. Received a tuple of length "
+                    f"{len(item)}."
+                )
+        else:
+            name = str(_FakePyTree[item])
+
+            class X(PyTree):
+                leaftype = item
+                structure = None
+
+        X.__name__ = name
+        X.__qualname__ = name
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            X.__module__ = "builtins"
+        else:
+            X.__module__ = "jaxtyping"
+        return X
+
+
+try:
+    # new typeguard
+    _TypeCheckError = (TypeError, typeguard.TypeCheckError)
+except AttributeError:
+    # old typeguard
+    _TypeCheckError = TypeError
+
+
+# Can't do `class PyTree(Generic[_T]): ...` because we need to override the
+# 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.
+
+Annotations of the following sorts are supported:
+```python
+a: PyTree
+b: PyTree[LeafType]
+c: PyTree[LeafType, "T"]
+d: PyTree[LeafType, "S T"]
+e: PyTree[LeafType, "... T"]
+f: PyTree[LeafType, "T ..."]
+```
+
+These correspond to:
+
+a. A plain `PyTree` can be used an annotation, 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))
+
+b. `PyTree[LeafType]` denotes a PyTree all of whose leaves match `LeafType`. For
+    example, `PyTree[int]` or `PyTree[Union[str, Float32[Array, "b c"]]]`.
+
+c. A structure name can also be passed. In this case
+    `jax.tree_util.tree_structure(...)` will be called, and bound to the structure name.
+    This can be used to mark that multiple PyTrees all have the same structure:
+    ```python
+    def f(x: PyTree[int, "T"], y: PyTree[int, "T"]):
+        ...
+    ```
+    Structures are bound to names in the same way as array shape annotations, i.e.
+    within the thread-local dynamic context of a [`jaxtyping.jaxtyped`][] decorator.
+
+d. A composite structure can be declared. In this case the variable must have a PyTree
+    structure each to the composition of multiple previously-bound PyTree structures.
+    For example:
+    ```python
+    def f(x: PyTree[int, "T"], y: PyTree[int, "S"], z: PyTree[int, "S T"]):
+        ...
+
+    x = (1, 2)
+    y = {"key": 3}
+    z = {"key": (4, 5)}  # structure is the composition of the structures of `y` and `z`
+    f(x, y, z)
+    ```
+    When performing runtime type-checking, all the individual pieces must have already
+    been bound to structures, otherwise the composite structure check will throw an error.
+
+e. A structure can begin with a `...`, to denote that the lower levels of the PyTree
+    must match the declared structure, but the upper levels can be arbitrary. As in the
+    previous case, all named pieces must already have been seen and their structures
+    bound.
+
+f. A structure can end with a `...`, to denote that the PyTree must be a prefix of the
+    declared structure, but the lower levels can be arbitrary. As in the previous two
+    cases, all named pieces must already have been seen and their structures bound.
+"""  # noqa: E501

jaxtyping/_raise.py

@@ -0,0 +1,23 @@
+from typing import NoReturn
+
+
+def jaxtyping_raise(e) -> NoReturn:
+    """Raises `e`, whilst adding a tag that it should not be intercepted by
+    `TypeCheckError`. All `raise` statements from within `__instancecheck__` should use
+    this.
+    """
+    __tracebackhide__ = True
+    try:
+        raise e
+    except Exception as f:
+        f._jaxtyping_malformed = True
+        raise
+
+
+def jaxtyping_raise_from(e, e_base) -> NoReturn:
+    __tracebackhide__ = True
+    try:
+        raise e from e_base
+    except Exception as f:
+        f._jaxtyping_malformed = True
+        raise

jaxtyping/_storage.py

@@ -0,0 +1,124 @@
+# 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
+from typing import Any, Optional
+
+from ._raise import jaxtyping_raise
+
+
+_shape_storage = threading.local()
+
+
+def _has_shape_memo():
+    return hasattr(_shape_storage, "memo_stack") and len(_shape_storage.memo_stack) != 0
+
+
+def get_shape_memo():
+    if _has_shape_memo():
+        single_memo, variadic_memo, pytree_memo, arguments = _shape_storage.memo_stack[
+            -1
+        ]
+    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.
+        single_memo = {}
+        variadic_memo = {}
+        pytree_memo = {}
+        arguments = {}
+    return single_memo, variadic_memo, pytree_memo, arguments
+
+
+def set_shape_memo(single_memo, variadic_memo, pytree_memo, arg_memo) -> None:
+    if _has_shape_memo():
+        _shape_storage.memo_stack[-1] = (
+            single_memo,
+            variadic_memo,
+            pytree_memo,
+            arg_memo,
+        )
+
+
+def push_shape_memo(arguments: dict[str, Any]):
+    try:
+        memo_stack = _shape_storage.memo_stack
+    except AttributeError:
+        # Can't be done when `_stack_storage` is created for reasons I forget.
+        memo_stack = _shape_storage.memo_stack = []
+    memos = ({}, {}, {}, arguments.copy())
+    memo_stack.append(memos)
+    return memos
+
+
+def pop_shape_memo() -> None:
+    _shape_storage.memo_stack.pop()
+
+
+_treepath_storage = threading.local()
+
+
+def clear_treepath_memo() -> None:
+    _treepath_storage.value = None
+
+
+def set_treepath_memo(index: Optional[int], structure: str) -> None:
+    if hasattr(_treepath_storage, "value") and _treepath_storage.value is not None:
+        jaxtyping_raise(
+            ValueError(
+                "Cannot typecheck annotations of the form "
+                "`PyTree[PyTree[Shaped[Array, '?foo'], 'T'], 'S']` as it is ambiguous "
+                "which PyTree the `?` annotation refers to."
+            )
+        )
+    if index is None:
+        _treepath_storage.value = f"~~delete~~({structure}) "
+    else:
+        # Appears in error messages, so human-readable
+        _treepath_storage.value = f"(Leaf {index} in structure {structure}) "
+
+
+def get_treepath_memo() -> str:
+    if not hasattr(_treepath_storage, "value") or _treepath_storage.value is None:
+        jaxtyping_raise(
+            ValueError(
+                "Cannot use `?` annotations, e.g. `Shaped[Array, '?foo']`, except "
+                "when contained with structured `PyTree` annotations, e.g. "
+                "`PyTree[Shaped[Array, '?foo'], 'T']`."
+            )
+        )
+    return _treepath_storage.value
+
+
+_treeflatten_storage = threading.local()
+
+
+def clear_treeflatten_memo() -> None:
+    _treeflatten_storage.value = False
+
+
+def set_treeflatten_memo():
+    _treeflatten_storage.value = True
+
+
+def get_treeflatten_memo():
+    try:
+        return _treeflatten_storage.value
+    except AttributeError:
+        return False

jaxtyping/array_types.py

@@ -1,522 +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 enum
-import functools as ft
-import typing
-from typing import Any, Dict, List, NoReturn, Optional, Tuple, TYPE_CHECKING, Union
-from typing_extensions import Literal
-
-import numpy as np
-
-from .decorator import storage
-
-
-_array_name_format = "dtype_and_shape"
-
-
-def get_array_name_format():
-    return _array_name_format
-
-
-def set_array_name_format(value):
-    global _array_name_format
-    _array_name_format = value
-
-
-_any_dtype = object()
-
-
-_anonymous_dim = object()
-_anonymous_variadic_dim = object()
-
-
-class _DimType(enum.Enum):
-    named = enum.auto()
-    fixed = enum.auto()
-    symbolic = enum.auto()
-
-
-class _NamedDim:
-    def __init__(self, name, broadcastable):
-        self.name = name
-        self.broadcastable = broadcastable
-
-
-class _NamedVariadicDim:
-    def __init__(self, name, broadcastable):
-        self.name = name
-        self.broadcastable = broadcastable
-
-
-class _FixedDim:
-    def __init__(self, size, broadcastable):
-        self.size = size
-        self.broadcastable = broadcastable
-
-
-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, _SymbolicDim]
-
-
-def _check_dims(
-    cls_dims: List[_AbstractDim],
-    obj_shape: 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:
-            pass
-        elif cls_dim.broadcastable and obj_size == 1:
-            pass
-        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 = single_memo[cls_dim.name]
-            except KeyError:
-                single_memo[cls_dim.name] = obj_size
-            else:
-                if cls_size != obj_size:
-                    return False
-    return True
-
-
-class _MetaAbstractArray(type):
-    def __instancecheck__(cls, obj):
-        if not isinstance(obj, cls.array_type):
-            return False
-
-        if 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 cls.dtypes is not _any_dtype and dtype not in cls.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.
-            single_memo = {}
-            variadic_memo = {}
-            variadic_broadcast_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 no_temp_memo:
-                storage.memo_stack[-1] = (
-                    single_memo,
-                    variadic_memo,
-                    variadic_broadcast_memo,
-                )
-            return True
-        else:
-            return False
-
-    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, single_memo)
-        else:
-            if obj.ndim < len(cls.dims) - 1:
-                return False
-            i = cls.index_variadic
-            j = -(len(cls.dims) - i - 1)
-            if j == 0:
-                j = None
-            if not _check_dims(cls.dims[:i], obj.shape[:i], single_memo):
-                return False
-            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 _anonymous_variadic_dim:
-                return True
-            else:
-                assert type(variadic_dim) is _NamedVariadicDim
-                variadic_name = variadic_dim.name
-                try:
-                    if variadic_dim.broadcastable:
-                        variadic_shapes = variadic_broadcast_memo[variadic_name]
-                    else:
-                        variadic_shape = variadic_memo[variadic_name]
-                except KeyError:
-                    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_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]
-            assert False
-
-
-class AbstractArray(metaclass=_MetaAbstractArray):
-    array_type: Any
-    dtypes: List[str]
-    dims: List[_AbstractDimOrVariadicDim]
-    index_variadic: Optional[int]
-
-
-class _MetaAbstractDtype(type):
-    def __instancecheck__(cls, obj: Any) -> NoReturn:
-        raise RuntimeError(
-            f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to "
-            "check just the dtype of an array, then use "
-            f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
-        )
-
-    @ft.lru_cache(maxsize=None)
-    def __getitem__(cls, item: Tuple[Any, str]) -> _MetaAbstractArray:
-        if not isinstance(item, tuple) or len(item) != 2:
-            raise ValueError(
-                "As of jaxtyping v0.2.0, type annotations must now include an explicit "
-                "array type. For example `jaxtyping.Float32[jnp.ndarray, 'foo bar']`."
-            )
-        array_type, dim_str = item
-        del item
-        if not isinstance(dim_str, str):
-            raise ValueError(
-                "Shape specification must be a string. Axes should be separated with "
-                "spaces."
-            )
-        dims = []
-        index_variadic = None
-        for index, elem in enumerate(dim_str.split()):
-            if "," in elem:
-                # Common mistake
-                raise ValueError(
-                    "Dimensions should be separated with spaces, not commas"
-                )
-            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:]
-                    else:
-                        break
-                try:
-                    elem = int(elem)
-                except ValueError:
-                    if len(elem) == 0 or elem.isidentifier():
-                        dim_type = _DimType.named
-                    else:
-                        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)
-        # In python 3.8, e.g., typing.Union lacks `__name__`.
-        try:
-            type_str = array_type.__name__
-        except AttributeError:
-            type_str = repr(array_type)
-        if _array_name_format == "dtype_and_shape":
-            name = f"{cls.__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")
-        out = _MetaAbstractArray(
-            name,
-            (AbstractArray,),
-            dict(
-                array_type=array_type,
-                dtypes=cls.dtypes,
-                dims=dims,
-                index_variadic=index_variadic,
-            ),
-        )
-        if getattr(typing, "GENERATING_DOCUMENTATION", False):
-            out.__module__ = "builtins"
-        else:
-            out.__module__ = "jaxtyping"
-        return out
-
-
-class AbstractDtype(metaclass=_MetaAbstractDtype):
-    dtypes: Union[Literal[_any_dtype], List[str]]
-
-    def __init__(self, *args, **kwargs):
-        raise RuntimeError(
-            "AbstractDtype cannot be instantiated. Perhaps you wrote e.g. "
-            '`Float32("shape")` when you mean `Float32[jnp.ndarray, "shape"]`?'
-        )
-
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-
-        dtypes: Union[Literal[_any_dtype], str, List[str]] = cls.dtypes
-        if isinstance(dtypes, str):
-            dtypes = [dtypes]
-        cls.dtypes = dtypes
-
-
-if TYPE_CHECKING:
-    # Note that `from typing_extensions import Annotated; ... = 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,
-        Annotated as Bool,
-        Annotated as Complex,
-        Annotated as Complex64,
-        Annotated as Complex128,
-        Annotated as Float,
-        Annotated as Float16,
-        Annotated as Float32,
-        Annotated as Float64,
-        Annotated as Inexact,
-        Annotated as Int,
-        Annotated as Int8,
-        Annotated as Int16,
-        Annotated as Int32,
-        Annotated as Int64,
-        Annotated as Integer,
-        Annotated as Num,
-        Annotated as Shaped,
-        Annotated as UInt,
-        Annotated as UInt8,
-        Annotated as UInt16,
-        Annotated as UInt32,
-        Annotated as UInt64,
-    )
-else:
-    _bool = "bool_"
-    _uint8 = "uint8"
-    _uint16 = "uint16"
-    _uint32 = "uint32"
-    _uint64 = "uint64"
-    _int8 = "int8"
-    _int16 = "int16"
-    _int32 = "int32"
-    _int64 = "int64"
-    _bfloat16 = "bfloat16"
-    _float16 = "float16"
-    _float32 = "float32"
-    _float64 = "float64"
-    _complex64 = "complex64"
-    _complex128 = "complex128"
-
-    def _make_dtype(_dtypes, name):
-        class _Cls(AbstractDtype):
-            dtypes = _dtypes
-
-        _Cls.__name__ = name
-        _Cls.__qualname__ = name
-        if getattr(typing, "GENERATING_DOCUMENTATION", False):
-            _Cls.__module__ = "builtins"
-        else:
-            _Cls.__module__ = "jaxtyping"
-        return _Cls
-
-    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")
-
-    uints = [_uint8, _uint16, _uint32, _uint64]
-    ints = [_int8, _int16, _int32, _int64]
-    floats = [_bfloat16, _float16, _float32, _float64]
-    complexes = [_complex64, _complex128]
-
-    # We match NumPy's type hierarachy in what types to provide. See the diagram at
-    # https://numpy.org/doc/stable/reference/arrays.scalars.html#scalars
-
-    Bool = _make_dtype(_bool, "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")

jaxtyping/decorator.py

@@ -1,76 +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 dataclasses
-import functools as ft
-import inspect
-import threading
-
-
-storage = threading.local()
-
-
-class _Jaxtyped:
-    def __init__(self, fn):
-        self.fn = fn
-
-    def __get__(self, instance, owner):
-        return ft.wraps(self.fn)(_Jaxtyped(self.fn.__get__(instance, owner)))
-
-    def __call__(self, *args, **kwargs):
-        try:
-            memo_stack = storage.memo_stack
-        except AttributeError:
-            memo_stack = storage.memo_stack = []
-        memo_stack.append(({}, {}, {}))
-        try:
-            return self.fn(*args, **kwargs)
-        finally:
-            memo_stack.pop()
-
-
-def jaxtyped(fn):
-    if 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"
-            )
-    else:
-        return ft.wraps(fn)(_Jaxtyped(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

@@ -1,299 +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.
-
-# This source code is adapted from typeguard:
-# https://github.com/agronholm/typeguard/blob/0dd7f7510b7c694e66a0d17d1d58d185125bad5d/src/typeguard/importhook.py
-#
-# Copied and adapted in compliance with the terms of typeguard's MIT license.
-# The original license is reproduced here.
-#
-# ---------
-#
-# This is the MIT license: http://www.opensource.org/licenses/mit-license.php
-#
-# Copyright (c) Alex Grönholm
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy of this
-# software and associated documentation files (the "Software"), to deal in the Software
-# without restriction, including without limitation the rights to use, copy, modify,
-# merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
-# permit persons to whom the Software is furnished to do so, subject to the following
-# conditions:
-#
-# The above copyright notice and this permission notice shall be included in all copies
-# or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
-# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
-# PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
-# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
-# CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE
-# OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
-#
-# ---------
-
-
-import ast
-import sys
-from importlib.abc import MetaPathFinder
-from importlib.machinery import SourceFileLoader
-from importlib.util import cache_from_source, decode_source
-from inspect import isclass
-from typing import Iterable, List, Optional, Tuple
-from unittest.mock import patch
-
-
-# The name of this function is magical
-def _call_with_frames_removed(f, *args, **kwargs):
-    return f(*args, **kwargs)
-
-
-def _optimized_cache_from_source(path, debug_override=None):
-    # 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.
-    return cache_from_source(path, debug_override, optimization="jaxtyping3")
-
-
-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
-
-
-class _JaxtypingTransformer(ast.NodeVisitor):
-    def __init__(self, *, typechecker) -> None:
-        self._parents: List[ast.AST] = []
-        self._typechecker = typechecker
-
-    def visit_Module(self, node: ast.Module):
-        # Insert "import typeguard; import jaxtping" after any "from __future__ ..."
-        # imports
-        for i, child in enumerate(node.body):
-            if isinstance(child, ast.ImportFrom) and child.module == "__future__":
-                continue
-            elif isinstance(child, ast.Expr) and isinstance(child.value, ast.Str):
-                continue  # module docstring
-            else:
-                node.body.insert(i, ast.Import(names=[ast.alias("jaxtyping", None)]))
-                if self._typechecker is not None:
-                    typechecker_module, _ = self._typechecker
-                    node.body.insert(
-                        i, ast.Import(names=[ast.alias(typechecker_module, None)])
-                    )
-                break
-
-        self._parents.append(node)
-        self.generic_visit(node)
-        self._parents.pop()
-        return node
-
-    def visit_ClassDef(self, node: ast.ClassDef):
-        func = _dot_lookup("jaxtyping", "decorator", "_jaxtyped_typechecker")
-        if self._typechecker is None:
-            args = [ast.Constant(None)]
-        else:
-            args = [_dot_lookup(*self._typechecker)]
-        node.decorator_list.append(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 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.
-            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.
-                node.decorator_list.append(_dot_lookup(*self._typechecker))
-        self._parents.append(node)
-        self.generic_visit(node)
-        self._parents.pop()
-        return node
-
-
-class _JaxtypingLoader(SourceFileLoader):
-    def __init__(self, *args, typechecker, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._typechecker = typechecker
-
-    def source_to_code(self, data, path, *, _optimize=-1):
-        source = decode_source(data)
-        tree = _call_with_frames_removed(
-            compile,
-            source,
-            path,
-            "exec",
-            ast.PyCF_ONLY_AST,
-            dont_inherit=True,
-            optimize=_optimize,
-        )
-        tree = _JaxtypingTransformer(typechecker=self._typechecker).visit(tree)
-        ast.fix_missing_locations(tree)
-        return _call_with_frames_removed(
-            compile, tree, path, "exec", dont_inherit=True, optimize=_optimize
-        )
-
-    def exec_module(self, module):
-        # Use a custom optimization marker - the import lock should make this monkey
-        # patch safe
-        with patch(
-            "importlib._bootstrap_external.cache_from_source",
-            _optimized_cache_from_source,
-        ):
-            return super().exec_module(module)
-
-
-class _JaxtypingFinder(MetaPathFinder):
-    """Wraps another path finder and instruments the module with `@jaxtyped` and
-    `@typechecked` if `should_instrument()` returns `True`.
-
-    Should not be used directly, but rather via `install_import_hook`.
-    """
-
-    def __init__(self, modules, original_pathfinder, typechecker):
-        self.modules = modules
-        self._original_pathfinder = original_pathfinder
-        self._typechecker = typechecker
-
-    def find_spec(self, fullname, path=None, target=None):
-        if self.should_instrument(fullname):
-            spec = self._original_pathfinder.find_spec(fullname, path, target)
-            if spec is not None and isinstance(spec.loader, SourceFileLoader):
-                spec.loader = _JaxtypingLoader(
-                    spec.loader.name, spec.loader.path, typechecker=self._typechecker
-                )
-                return spec
-
-        return None
-
-    def should_instrument(self, module_name: str) -> bool:
-        """Determine whether the module with the given name should be instrumented.
-
-        **Arguments:**
-
-        - `module_name`: the full name of the module that is about to be imported
-            (e.g. ``xyz.abc``)
-        """
-        for module in self.modules:
-            if module_name == module or module_name.startswith(module + "."):
-                return True
-
-        return False
-
-
-class ImportHookManager:
-    def __init__(self, hook: MetaPathFinder):
-        self.hook = hook
-
-    def __enter__(self):
-        pass
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.uninstall()
-
-    def uninstall(self):
-        try:
-            sys.meta_path.remove(self.hook)
-        except ValueError:
-            pass  # already removed
-
-
-# Deliberately no default for `typechecker` so that folks must opt-in to not having
-# a typechecker.
-def install_import_hook(
-    modules: Iterable[str], typechecker: Optional[Tuple[str, str]]
-) -> ImportHookManager:
-    """Automatically apply `@jaxtyped`, and optionally a type checker, to all classes
-    and functions.
-
-    It will only be applied to modules loaded **after** this hook has been installed.
-
-    **Arguments:**:
-
-    - `packages`: the names of the modules in which to automatically apply `@jaxtyped`
-        and `@typechecked`.
-    - `typechecker`: the module and function of the typechecker you want to use, as a
-        2-tuple of strings. For example `typechecker=("typeguard", "typechecked")` or
-        `typechecker=("beartype", "beartype")`. You may pass `typechecker=None` if you
-        do not want to automatically decorate with a typechecker as well.
-
-    If the function already has any decorators on it, then both the `@jaxtyped` and the
-    typechecker decorators will go at the bottom of the decorator list, e.g.
-    ```python
-    @some_other_decorator
-    @jaxtyped
-    @beartype.beartype
-    def foo(...): ...
-    ```
-
-    **Returns:**
-
-    A context manager that uninstalls the hook on exit, or when you call `.uninstall()`.
-
-    **Example:**
-
-    Typically you should apply this import hook at the entry point for your own scripts:
-    ```python
-    # entry_point.py
-    from jaxtyped import install_import_hook
-    with install_import_hook("main", ("beartype", "beartype"))
-        import main
-    ...  # do whatever you're doing
-
-    # main.py
-    from jaxtyped import f32
-
-    def f(x: f32["b c"]):
-        pass
-    ```
-    Which as you can see means you never to import `@jaxtyped`, nor do you need to
-    import the typechecker directly (e.g. `beartype.beartype` or
-    `typeguard.typechecked`).
-    """
-    if isinstance(modules, str):
-        modules = [modules]
-
-    for i, finder in enumerate(sys.meta_path):
-        if (
-            isclass(finder)
-            and finder.__name__ == "PathFinder"
-            and hasattr(finder, "find_spec")
-        ):
-            break
-    else:
-        raise RuntimeError("Cannot find a PathFinder in sys.meta_path")
-
-    hook = _JaxtypingFinder(modules, finder, typechecker)
-    sys.meta_path.insert(0, hook)
-    return ImportHookManager(hook)

jaxtyping/pytree_type.py

@@ -1,104 +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 functools as ft
-import typing
-from typing import Generic, TypeVar
-
-import jax.tree_util as jtu
-import typeguard
-
-
-_T = TypeVar("_T")
-
-
-class _FakePyTree(Generic[_T]):
-    pass
-
-
-_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__
-# "types", e.g. we get types.PyTree[int].
-
-
-class _MetaPyTree(type):
-    def __call__(self, *args, **kwargs):
-        raise RuntimeError("PyTree cannot be instantiated")
-
-    def __instancecheck__(cls, obj):
-        return True
-
-    @ft.lru_cache(maxsize=None)
-    def __getitem__(cls, item):
-        name = str(_FakePyTree[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):
-    def __call__(self, *args, **kwargs):
-        raise RuntimeError("PyTree cannot be instantiated")
-
-    def __instancecheck__(cls, obj):
-        # We could use `isinstance` here but that would fail for more complicated
-        # types, e.g. PyTree[Tuple[int]]. So at least internally we make a particular
-        # choice of typechecker.
-        #
-        # Deliberately not using @jaxtyped so that we share the same `memo` as whatever
-        # dynamic context we're currently in.
-        @typeguard.typechecked
-        def accepts_leaftype(x: cls.leaftype):
-            pass
-
-        def is_leaftype(x):
-            try:
-                accepts_leaftype(x)
-            except _TypeCheckError:
-                return False
-            else:
-                return True
-
-        leaves = jtu.tree_leaves(obj, is_leaf=is_leaftype)
-        return all(map(is_leaftype, leaves))
-
-
-# Can't do `class PyTree(Generic[_T]): ...` because we need to override the
-# 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"

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.24"
+description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
+readme = "README.md"
+requires-python ="~=3.9"
+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,<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,100 +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.
-# For typing_extensions, we choose versions that match
-# https://github.com/explosion/confection/blob/main/setup.cfg#L33 used in colab
-
-install_requires = [
-    "numpy>=1.20.0",
-    "typeguard>=2.13.3",
-    "typing_extensions>=3.7.4.1",
-]
-
-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],
-    include_package_data=True,
-)

test/conftest.py

@@ -41,6 +41,29 @@ def typecheck(request):
     return request.param
 
 
+@pytest.fixture(params=(False, True))
+def jaxtyp(request):
+    import jaxtyping
+
+    if request.param:
+        # New-style
+        # @jaxtyping.jaxtyped(typechecker=typechecker)
+        # def f(...)
+        return lambda typechecker: jaxtyping.jaxtyped(typechecker=typechecker)
+    else:
+        # Old-style
+        # @jaxtyping.jaxtyped
+        # @typechecker
+        # def f(...)
+        def impl(typechecker):
+            def decorator(fn):
+                return jaxtyping.jaxtyped(typechecker(fn))
+
+            return decorator
+
+        return impl
+
+
 @pytest.fixture()
 def getkey():
     def _getkey():
@@ -48,3 +71,13 @@ def getkey():
         return jr.PRNGKey(random.randint(0, 2**31 - 1))
 
     return _getkey
+
+
+@pytest.fixture(scope="module")
+def beartype_or_skip():
+    yield pytest.importorskip("beartype")
+
+
+@pytest.fixture(scope="module")
+def typeguard_or_skip():
+    yield pytest.importorskip("typeguard")

test/import_hook_tester.py

@@ -0,0 +1,234 @@
+# 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 helpers import ParamError, ReturnError
+
+import jaxtyping
+from jaxtyping import Float32, Int
+
+
+#
+# Test that functions get checked
+#
+
+
+def g(x: Float32[jnp.ndarray, " b"]):
+    pass
+
+
+g(jnp.array([1.0]))
+with pytest.raises(ParamError):
+    g(jnp.array(1))
+
+
+#
+# Test that Equinox modules get checked
+#
+
+
+# Dataclass `__init__`, no converter
+class Mod1(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+Mod1(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod1(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod1(1, jnp.array(1.0))
+
+
+# Dataclass `__init__`, converter
+class Mod2(eqx.Module):
+    a: jnp.ndarray = eqx.field(converter=jnp.asarray)
+
+
+Mod2(1)  # This will fail unless we run typechecking after conversion
+
+
+class BadMod2(eqx.Module):
+    a: jnp.ndarray = eqx.field(converter=lambda x: x)
+
+
+with pytest.raises(ParamError):
+    BadMod2(1)
+with pytest.raises(ParamError):
+    BadMod2("asdf")
+
+
+# Custom `__init__`, no converter
+class Mod3(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+    def __init__(self, foo: str, bar: Float32[jnp.ndarray, " a"]):
+        self.foo = int(foo)
+        self.bar = bar
+
+
+Mod3("1", jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod3(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod3("1", jnp.array(1.0))
+
+
+# Custom `__init__`, converter
+class Mod4(eqx.Module):
+    a: Int[jnp.ndarray, ""] = eqx.field(converter=jnp.asarray)
+
+    def __init__(self, a: str):
+        self.a = int(a)
+
+
+Mod4("1")  # This will fail unless we run typechecking after conversion
+
+
+# Custom `__post_init__`, no converter
+class Mod5(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+    def __post_init__(self):
+        pass
+
+
+Mod5(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod5(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod5(1, jnp.array(1.0))
+
+
+# Dataclass `__init__`, converter
+class Mod6(eqx.Module):
+    a: jnp.ndarray = eqx.field(converter=jnp.asarray)
+
+    def __post_init__(self):
+        pass
+
+
+Mod6(1)  # This will fail unless we run typechecking after conversion
+
+
+#
+# Test that dataclasses get checked
+#
+
+
+@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 that methods get checked
+#
+
+
+class N(eqx.Module):
+    a: jnp.ndarray
+
+    def __init__(self, foo: str):
+        self.a = jnp.array(1)
+
+    def foo(self, x: jnp.ndarray):
+        pass
+
+    def bar(self) -> jnp.ndarray:
+        return self.a
+
+
+n = N("hi")
+with pytest.raises(ParamError):
+    N(123)
+with pytest.raises(ParamError):
+    n.foo("not_an_array_either")
+bad_n = eqx.tree_at(lambda x: x.a, n, "not_an_array")
+with pytest.raises(ReturnError):
+    bad_n.bar()
+
+
+#
+# Test that we don't get called in `super()`.
+#
+
+
+called = False
+
+
+class Base(eqx.Module):
+    x: int
+
+    def __init__(self):
+        self.x = "not an int"
+        global called
+        assert not called
+        called = True
+
+
+class Derived(Base):
+    def __init__(self):
+        assert not called
+        super().__init__()
+        assert called
+        self.x = 2
+
+
+Derived()
+
+
+#
+# Test that stringified type annotations work
+
+
+class Foo:
+    pass
+
+
+class Bar(eqx.Module):
+    x: type[Foo]
+    y: "type[Foo]"
+    # Note that this is the *only* kind of partially-stringified type annotation that
+    # is supported. This is for compatibility with older Equinox versions.
+    z: type["Foo"]
+
+
+Bar(Foo, Foo, Foo)
+
+with pytest.raises(ParamError):
+    Bar(1, Foo, Foo)
+
+# Record that we've finished our checks successfully
+
+jaxtyping._test_import_hook_counter += 1

test/import_hook_tester_beartype.py

@@ -1,47 +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 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))

test/import_hook_tester_broken_checker.py

@@ -1,47 +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 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))

test/import_hook_tester_transitive/__init__.py

@@ -1,20 +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.
-
-from . import another_file

test/import_hook_tester_transitive/another_file.py

@@ -1,47 +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 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))

test/import_hook_tester_typeguard.py

@@ -1,47 +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 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))

test/requirements.txt

@@ -1,4 +1,8 @@
-equinox>=0.5.3
-pytest>=7.0.1
-beartype>=0.10.4
-typeguard>=2.13.3
+beartype
+cloudpickle
+equinox
+IPython
+jaxlib
+pytest
+tensorflow
+typeguard<3

test/test_array.py

@@ -17,18 +17,35 @@
 # 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 as dc
+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
+import torch
 
-from jaxtyping import AbstractDtype, Array, Float, Float32, jaxtyped, Shaped
+from jaxtyping import (
+    AbstractDtype,
+    Array,
+    ArrayLike,
+    Bool,
+    Float,
+    Float32,
+    jaxtyped,
+    Key,
+    PRNGKeyArray,
+    Scalar,
+    Shaped,
+)
 
 from .helpers import ParamError, ReturnError
 
 
-def test_basic(typecheck):
-    @jaxtyped
-    @typecheck
+def test_basic(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "..."]):
         pass
 
@@ -67,16 +84,14 @@ def test_dtypes():
             assert key == val.__name__
 
 
-def test_return(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_return(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]:
         return jnp.transpose(x)
 
     g(jr.normal(getkey(), (3, 4)))
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float[Array, "b c"]) -> Float[Array, "b c"]:
         return jnp.transpose(x)
 
@@ -84,9 +99,8 @@ def test_return(typecheck, getkey):
         h(jr.normal(getkey(), (3, 4)))
 
 
-def test_two_args(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_two_args(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]):
         return x @ y
 
@@ -94,8 +108,7 @@ def test_two_args(typecheck, getkey):
     with pytest.raises(ParamError):
         g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]) -> Shaped[Array, "b d"]:
         return x @ y
 
@@ -104,9 +117,8 @@ def test_two_args(typecheck, getkey):
         h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
 
 
-def test_any_dtype(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_any_dtype(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
         return x
 
@@ -121,14 +133,12 @@ def test_any_dtype(typecheck, getkey):
         g(jr.normal(getkey(), (1,)))
 
 
-def test_nested_jaxtyped(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_nested_jaxtyped(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"], transpose: bool) -> Float32[Array, "c b"]:
         return h(x, transpose)
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float32[Array, "c b"], transpose: bool) -> Float32[Array, "b c"]:
         if transpose:
             return jnp.transpose(x)
@@ -142,9 +152,8 @@ def test_nested_jaxtyped(typecheck, getkey):
         g(jr.normal(getkey(), (2, 3)), False)
 
 
-def test_nested_nojaxtyped(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_nested_nojaxtyped(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"]):
         return h(x)
 
@@ -156,9 +165,8 @@ def test_nested_nojaxtyped(typecheck, getkey):
         g(jr.normal(getkey(), (2, 3)))
 
 
-def test_isinstance(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_isinstance(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"]) -> Float32[Array, " z"]:
         y = jnp.transpose(x)
         assert isinstance(y, Float32[Array, "c b"])
@@ -172,9 +180,8 @@ def test_isinstance(typecheck, getkey):
     g(jr.normal(getkey(), (2, 3)))
 
 
-def test_fixed(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_fixed(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(
         x: Float32[Array, "4 5 foo"], y: Float32[Array, " foo"]
     ) -> Float32[Array, "4 5"]:
@@ -189,9 +196,8 @@ def test_fixed(typecheck, getkey):
         g(c, b)
 
 
-def test_anonymous(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_anonymous(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "foo _"], y: Float32[Array, " _"]):
         pass
 
@@ -200,9 +206,8 @@ def test_anonymous(typecheck, getkey):
     g(a, b)
 
 
-def test_named_variadic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_named_variadic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(
         x: Float32[Array, "*batch foo"],
         y: Float32[Array, " *batch"],
@@ -225,8 +230,7 @@ def test_named_variadic(typecheck, getkey):
     with pytest.raises(ParamError):
         g(a2, b1, c)
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float32[Array, " foo *batch"], y: Float32[Array, " foo *batch bar"]):
         pass
 
@@ -240,9 +244,8 @@ def test_named_variadic(typecheck, getkey):
         h(b, c)
 
 
-def test_anonymous_variadic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_anonymous_variadic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "... foo"], y: Float32[Array, " foo"]):
         pass
 
@@ -262,9 +265,8 @@ def test_anonymous_variadic(typecheck, getkey):
         g(a3, c)
 
 
-def test_broadcast_fixed(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_fixed(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "#4"]):
         pass
 
@@ -275,9 +277,8 @@ def test_broadcast_fixed(typecheck, getkey):
         g(jr.normal(getkey(), (3,)))
 
 
-def test_broadcast_named(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_named(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, " #foo"], y: Float32[Array, " #foo"]):
         pass
 
@@ -299,9 +300,8 @@ def test_broadcast_named(typecheck, getkey):
         g(b, a)
 
 
-def test_broadcast_variadic_named(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_variadic_named(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, " *#foo"], y: Float32[Array, " *#foo"]):
         pass
 
@@ -358,29 +358,71 @@ def test_broadcast_variadic_named(typecheck, getkey):
         g(o, a)
 
 
+def test_variadic_mixed_broadcast(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(x: Float[Array, " *foo"], y: Float[Array, " #*foo"]):
+        pass
+
+    a = jr.normal(getkey(), (3, 4))
+    b = jr.normal(getkey(), (5,))
+    with pytest.raises(ParamError):
+        f(a, b)
+
+    c = jr.normal(getkey(), (7, 3, 2))
+    d = jr.normal(getkey(), (1, 2))
+    f(c, d)
+
+
+def test_variadic_mixed_broadcast2(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(x: Float[Array, " *#foo"], y: Float[Array, " *foo"]):
+        pass
+
+    a = jr.normal(getkey(), (3, 4))
+    b = jr.normal(getkey(), (5,))
+    with pytest.raises(ParamError):
+        f(a, b)
+
+    c = jr.normal(getkey(), (1, 2))
+    d = jr.normal(getkey(), (7, 3, 2))
+    f(c, d)
+
+
+def test_variadic_mixed_broadcast3(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(
+        x: Float[Array, "*B L D"],
+        *,
+        y: Float[Array, "*#B J d"],
+        z: Bool[Array, "*B L J"],
+    ) -> Float[Array, "*B L D"]:
+        return x
+
+    x = jr.normal(getkey(), (2, 7, 3, 2, 2))
+    y = jr.bernoulli(getkey(), shape=(2, 7, 3, 2, 2))
+    z = jr.normal(getkey(), (2, 7, 1, 2, 2))
+    f(x, y=z, z=y)
+
+
 def test_no_commas():
     with pytest.raises(ValueError):
         Float32[Array, "foo, bar"]
 
 
-def test_symbolic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_symbolic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
         return x[1:]
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
         return jnp.concatenate([x, x])
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def bad_make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
         return x
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def bad_cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
         return jnp.concatenate([x, x, x])
 
@@ -400,12 +442,241 @@ def test_symbolic(typecheck, getkey):
         bad_cat(x)
 
 
-def test_incomplete_symbolic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_incomplete_symbolic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def foo(x: Float32[Array, " 2*dim"]):
         pass
 
     x = jr.normal(getkey(), (4,))
     with pytest.raises(NameError):
         foo(x)
+
+
+def test_deferred_symbolic_good(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def foo(dim: int, fill: Float[Array, ""]) -> Float[Array, " {dim}"]:
+        return jnp.full((dim,), fill)
+
+    class A:
+        size = 5
+
+        @jaxtyp(typecheck)
+        def bar(self, fill: Float[Array, ""]) -> Float[Array, " {self.size}"]:
+            return jnp.full((self.size,), fill)
+
+    foo(3, jnp.array(0.0))
+    A().bar(jnp.array(0.0))
+
+
+def test_deferred_symbolic_bad(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def foo(dim: int, fill: Float[Array, ""]) -> Float[Array, " {dim-1}"]:
+        return jnp.full((dim,), fill)
+
+    class A:
+        size = 5
+
+        @jaxtyp(typecheck)
+        def bar(self, fill: Float[Array, ""]) -> Float[Array, " {self.size}-1"]:
+            return jnp.full((self.size,), fill)
+
+    with pytest.raises(ReturnError):
+        foo(3, jnp.array(0.0))
+
+    with pytest.raises(ReturnError):
+        A().bar(jnp.array(0.0))
+
+
+def test_deferred_symbolic_dataclass(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    @dc.dataclass
+    class A:
+        value: int
+        array: Float[Array, " {self.value}"]
+
+    A(3, jnp.zeros(3))
+
+    with pytest.raises(ParamError):
+        A(3, jnp.zeros(4))
+
+
+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.number, ""],
+        float,
+    }
+    assert set(get_args(floatlike2)) == {
+        Float[Array, ""],
+        Float[np.ndarray, ""],
+        Float[np.number, ""],
+        float,
+    }
+    assert set(get_args(floatlike3)) == {
+        Float32[Array, "4"],
+        Float32[np.ndarray, "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"],
+    }
+
+
+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(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PRNGKeyArray):
+        pass
+
+    f(jr.key(0))
+    f(jr.PRNGKey(0))
+
+    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_key_dtype(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f1(x: Key[Array, ""]):
+        pass
+
+    @jaxtyp(typecheck)
+    def f2(x: Key[Scalar, ""]):
+        pass
+
+    for f in (f1, f2):
+        f(jr.key(0))
+
+        with pytest.raises(ParamError):
+            f(jr.PRNGKey(0))
+        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(jaxtyp, 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"]
+
+    @jaxtyp(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))
+
+
+def test_scalar_variadic_dim():
+    assert Float[float, "..."] is float
+    assert Float[float, "#*shape"] is float
+
+    # This one is a bit weird -- it should really also assert that shape==(), but we
+    # don't implement that.
+    assert Float[float, "*shape"] is float
+
+
+def test_scalar_dtype_mismatch():
+    with pytest.raises(ValueError):
+        Float[bool, "..."]

test/test_decorator.py

@@ -1,13 +1,166 @@
-from jaxtyping import jaxtyped
+import abc
+
+import jax.random as jr
+import pytest
+
+from jaxtyping import Array, Float, jaxtyped
+
+from .helpers import ParamError, ReturnError
 
 
-class M:
+class M(metaclass=abc.ABCMeta):
+    @jaxtyped
+    def f(self):
+        ...
+
     @jaxtyped
     @classmethod
-    def f(cls):
+    def g1(cls):
         return 3
 
+    @classmethod
+    @jaxtyped
+    def g2(cls):
+        return 4
 
-# Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod
-def test_decorator():
-    assert M.f() == 3
+    @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
+
+
+def test_context(getkey):
+    a = jr.normal(getkey(), (3, 4))
+    b = jr.normal(getkey(), (5,))
+    with jaxtyped("context"):
+        assert isinstance(a, Float[Array, "foo bar"])
+        assert not isinstance(b, Float[Array, "foo"])
+    assert isinstance(a, Float[Array, "foo bar"])
+    assert isinstance(b, Float[Array, "foo"])
+
+
+def test_varargs(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(*args):
+        pass
+
+    f(1, 2)
+
+
+def test_varkwargs(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(**kwargs):
+        pass
+
+    f(a=1, b=2)
+
+
+def test_defaults(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: int, y=1):
+        pass
+
+    f(1)
+
+
+class _GlobalFoo:
+    pass
+
+
+def test_global_stringified_annotation(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: "_GlobalFoo") -> "_GlobalFoo":
+        return x
+
+    f(_GlobalFoo())
+
+    @jaxtyp(typecheck)
+    def g(x: int) -> "_GlobalFoo":
+        return x
+
+    @jaxtyp(typecheck)
+    def h(x: "_GlobalFoo") -> int:
+        return x
+
+    with pytest.raises(ReturnError):
+        g(1)
+
+    with pytest.raises(ParamError):
+        h(1)
+
+
+# This test does not use `jaxtyp(typecheck)` because typeguard does some evil stack
+# frame introspection to try and grab local variables.
+def test_local_stringified_annotation(typecheck):
+    class LocalFoo:
+        pass
+
+    @jaxtyped(typechecker=typecheck)
+    def f(x: "LocalFoo") -> "LocalFoo":
+        return x
+
+    f(LocalFoo())
+
+    @jaxtyped
+    @typecheck
+    def g(x: "LocalFoo") -> "LocalFoo":
+        return x
+
+    g(LocalFoo())
+
+    # We don't check that errors are raised if it goes wrong, since we can't usually
+    # resolve local type annotations at runtime. Best we can hope for is not to raise
+    # a spurious error about not being able to find the type.

test/test_import_hook.py

@@ -17,43 +17,105 @@
 # 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
+import importlib.metadata
+import pathlib
+import shutil
+import sys
+import tempfile
+
 import pytest
 
-from jaxtyping import install_import_hook
+import jaxtyping
 
 
-def test_import_hook_typeguard():
-    hook = install_import_hook(
-        "test.import_hook_tester_typeguard", ("typeguard", "typechecked")
-    )
-    with hook:
-        from . import import_hook_tester_typeguard  # noqa: F401
+_here = pathlib.Path(__file__).parent
 
 
-def test_import_hook_beartype():
+try:
+    typeguard_version = importlib.metadata.version("typeguard")
+except Exception as e:
+    raise ImportError("Could not find typeguard version") from e
+else:
     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
-
-
-def test_import_hook_transitive():
-    hook = install_import_hook(
-        "test.import_hook_tester_transitive", ("typeguard", "typechecked")
+        major, _, _ = typeguard_version.split(".")
+        major = int(major)
+    except Exception as e:
+        raise ImportError(
+            f"Unexpected typeguard version {typeguard_version}; not formatted as "
+            "`major.minor.patch`"
+        ) from e
+if major != 2:
+    raise ImportError(
+        "jaxtyping's tests required typeguard version 2. (Versions 3 and 4 are both "
+        "known to have bugs.)"
     )
-    with hook:
-        from . import import_hook_tester_transitive  # noqa: F401
 
 
-def test_import_hook_broken_checker():
-    hook = install_import_hook(
-        "test.import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist")
-    )
-    with hook, pytest.raises(AttributeError):
-        from . import import_hook_tester_broken_checker  # noqa: F401
+assert not hasattr(jaxtyping, "_test_import_hook_counter")
+jaxtyping._test_import_hook_counter = 0
+
+
+@pytest.fixture(scope="module")
+def importhook_tempdir():
+    with tempfile.TemporaryDirectory() as dir:
+        sys.path.append(dir)
+        dir = pathlib.Path(dir)
+        shutil.copyfile(_here / "helpers.py", dir / "helpers.py")
+        yield dir
+
+
+def _test_import_hook(importhook_tempdir, typechecker):
+    counter = jaxtyping._test_import_hook_counter
+    stem = f"import_hook_tester{counter}"
+    shutil.copyfile(_here / "import_hook_tester.py", importhook_tempdir / f"{stem}.py")
+
+    importlib.invalidate_caches()
+    with jaxtyping.install_import_hook(stem, typechecker):
+        importlib.import_module(stem)
+    assert counter + 1 == jaxtyping._test_import_hook_counter
+
+
+# Tests start below...
+
+
+def test_import_hook_typeguard(importhook_tempdir, typeguard_or_skip):
+    _test_import_hook(importhook_tempdir, "typeguard.typechecked")
+
+
+def test_import_hook_beartype(importhook_tempdir, beartype_or_skip):
+    _test_import_hook(importhook_tempdir, "beartype.beartype")
+
+
+def test_import_hook_beartype_full(importhook_tempdir, beartype_or_skip):
+    bearchecker = "beartype.beartype(conf=beartype.BeartypeConf(strategy=beartype.BeartypeStrategy.On))"  # noqa: E501
+    _test_import_hook(importhook_tempdir, bearchecker)
+
+
+def test_import_hook_typeguard_old(importhook_tempdir, typeguard_or_skip):
+    _test_import_hook(importhook_tempdir, ("typeguard", "typechecked"))
+
+
+def test_import_hook_beartype_old(importhook_tempdir, beartype_or_skip):
+    _test_import_hook(importhook_tempdir, ("beartype", "beartype"))
+
+
+def test_import_hook_broken_checker(importhook_tempdir):
+    with pytest.raises(AttributeError):
+        _test_import_hook(importhook_tempdir, "jaxtyping.does_not_exist")
+
+
+def test_import_hook_transitive(importhook_tempdir, typeguard_or_skip):
+    counter = jaxtyping._test_import_hook_counter
+    transitive_name = "jaxtyping_transitive_test"
+    transitive_dir = importhook_tempdir / transitive_name
+    transitive_dir.mkdir()
+    shutil.copyfile(_here / "import_hook_tester.py", transitive_dir / "tester.py")
+    with open(transitive_dir / "__init__.py", "w") as f:
+        f.write("from . import tester")
+        f.flush()
+
+    importlib.invalidate_caches()
+    with jaxtyping.install_import_hook(transitive_name, "typeguard.typechecked"):
+        importlib.import_module(transitive_name)
+    assert counter + 1 == jaxtyping._test_import_hook_counter

test/test_ipython_extension.py

@@ -0,0 +1,151 @@
+import pytest
+from IPython.testing.globalipapp import start_ipython
+
+from .helpers import ParamError
+
+
+@pytest.fixture(scope="session")
+def session_ip():
+    yield start_ipython()
+
+
+@pytest.fixture(scope="function")
+def ip(session_ip):
+    session_ip.run_cell(raw_cell="import jaxtyping")
+    session_ip.run_line_magic(magic_name="load_ext", line="jaxtyping")
+    session_ip.run_line_magic(
+        magic_name="jaxtyping.typechecker", line="typeguard.typechecked"
+    )
+    yield session_ip
+
+
+def test_that_ipython_works(ip):
+    ip.run_cell(raw_cell="x = 1").raise_error()
+    assert ip.user_global_ns["x"] == 1
+
+
+def test_function_beartype(ip):
+    ip.run_cell(
+        raw_cell="""
+    def f(x: int):
+        pass
+                """
+    ).raise_error()
+    ip.run_cell(raw_cell="f(1)").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='f("x")').raise_error()
+
+
+def test_function_none(ip):
+    ip.run_cell(
+        raw_cell="""
+    def f(a,b,c):
+        pass
+                """
+    ).raise_error()
+    ip.run_cell(raw_cell='f(1,2,"k")').raise_error()
+
+
+def test_function_jaxtyped(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import jax
+
+    def g(x: Float[Array, "1"]):
+        return x + 1
+
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='g("string")').raise_error()
+
+
+def test_function_jaxtyped_and_jitted(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import jax
+
+    @jax.jit
+    def g(x: Float[Array, "1"]):
+        return x + 1
+
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='g("string")').raise_error()
+
+
+def test_class_jaxtyped(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import equinox as eqx
+    import jax
+
+    class A(eqx.Module):
+        x: Float[Array, "2"]
+
+        def do_something(self, y: Int[Array, ""]):
+            return self.x + y
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(
+            raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))"
+        ).raise_error()
+
+
+def test_class_not_dataclass(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import equinox as eqx
+    import jax
+
+    class A:
+        def __init__(self, x):
+            self.x = x
+
+        def do_something(self, y):
+            return x + y
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
+    ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))").raise_error()

test/test_messages.py

@@ -0,0 +1,104 @@
+from typing import Any
+
+import equinox as eqx
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Array, Float, jaxtyped, PyTree, TypeCheckError
+
+
+def test_arg_localisation(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    def f(x: str, y: str, z: int):
+        pass
+
+    matches = [
+        "Type-check error whilst checking the parameters of f",
+        "The problem arose whilst typechecking argument 'z'.",
+        "Called with arguments: {'x': 'hi', 'y': 'bye', 'z': 'not-an-int'}",
+        r"Parameter annotations: \(x: str, y: str, z: int\).",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            f("hi", "bye", "not-an-int")
+
+    @jaxtyped(typechecker=typecheck)
+    def g(x: Float[Array, "a b"], y: Float[Array, "b c"]):
+        pass
+
+    x = jnp.zeros((2, 3))
+    y = jnp.zeros((4, 3))
+    matches = [
+        "Type-check error whilst checking the parameters of g",
+        "The problem arose whilst typechecking argument 'y'.",
+        r"Called with arguments: {'x': f32\[2,3\], 'y': f32\[4,3\]}",
+        (
+            r"Parameter annotations: \(x: Float\[Array, 'a b'\], y: "
+            r"Float\[Array, 'b c'\]\)."
+        ),
+        "The current values for each jaxtyping axis annotation are as follows.",
+        "a=2",
+        "b=3",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            g(x, y=y)
+
+
+def test_return(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    def f(x: PyTree[Any, " T"], y: PyTree[Any, " S"]) -> PyTree[Any, "T S"]:
+        return "foo"
+
+    x = (1, 2)
+    y = {"a": 1}
+    matches = [
+        "Type-check error whilst checking the return value of f",
+        r"Called with arguments: {'x': \(1, 2\), 'y': {'a': 1}}",
+        "Return value: 'foo'",
+        r"Return annotation: PyTree\[Any, \"T S\"\].",
+        (
+            "The current values for each jaxtyping PyTree structure annotation are as "
+            "follows."
+        ),
+        r"T=PyTreeDef\(\(\*, \*\)\)",
+        r"S=PyTreeDef\({'a': \*}\)",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            f(x, y=y)
+
+
+def test_dataclass_attribute(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    class M(eqx.Module):
+        x: Float[Array, " *foo"]
+        y: PyTree[Any, " T"]
+        z: int
+
+    x = jnp.zeros((2, 3))
+    y = (1, (3, 4))
+    z = "not-an-int"
+
+    matches = [
+        "Type-check error whilst checking the parameters of M",
+        "The problem arose whilst typechecking argument 'z'.",
+        (
+            r"Called with arguments: {'self': M\(\.\.\.\), 'x': f32\[2,3\], "
+            r"'y': \(1, \(3, 4\)\), 'z': 'not-an-int'}"
+        ),
+        (
+            r"Parameter annotations: \(self: Any, x: Float\[Array, '\*foo'\], "
+            r"y: PyTree\[Any, \"T\"\], z: int\)."
+        ),
+        "The current values for each jaxtyping axis annotation are as follows.",
+        r"foo=\(2, 3\)",
+        (
+            "The current values for each jaxtyping PyTree structure annotation are as "
+            "follows."
+        ),
+        r"T=PyTreeDef\(\(\*, \(\*, \*\)\)\)",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            M(x, y, z)

test/test_pytree.py

@@ -25,7 +25,8 @@ import jax.numpy as jnp
 import jax.random as jr
 import pytest
 
-from jaxtyping import Float, jaxtyped, PyTree
+import jaxtyping
+from jaxtyping import Array, Float, PyTree
 
 from .helpers import make_mlp, ParamError
 
@@ -92,9 +93,8 @@ def test_nested_pytrees(getkey, typecheck):
         g([1, 2, make_mlp()])
 
 
-def test_pytree_array(typecheck):
-    @jaxtyped
-    @typecheck
+def test_pytree_array(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def g(x: PyTree[Float[jnp.ndarray, "..."]]):
         pass
 
@@ -106,9 +106,8 @@ def test_pytree_array(typecheck):
         g(1.0)
 
 
-def test_pytree_shaped_array(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_pytree_shaped_array(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: PyTree[Float[jnp.ndarray, "b c"]]):
         pass
 
@@ -185,3 +184,161 @@ def test_pytree_namedtuple(typecheck):
                 y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)),
             )
         )
+
+
+def test_subclass_pytree():
+    x = PyTree
+    y = PyTree[int]
+    assert issubclass(x, PyTree)
+    assert issubclass(y, PyTree)
+    assert not issubclass(int, PyTree)
+
+
+def test_structure_match(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[int, " T"], y: PyTree[str, " T"]):
+        pass
+
+    f(1, "hi")
+    f((3, 4, {"a": 5}), ("a", "b", {"a": "c"}))
+
+    with pytest.raises(ParamError):
+        f(1, ("hi",))
+
+
+def test_structure_prefix(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[int, " T"], y: PyTree[str, "T ..."]):
+        pass
+
+    f(1, "hi")
+    f((3, 4, {"a": 5}), ("a", "b", {"a": "c"}))
+    f(1, ("hi",))
+    f((1, 2), ({"a": "hi"}, {"a": "bye"}))
+    f((1, 2), ({"a": "hi"}, {"not-a": "bye"}))
+
+    with pytest.raises(ParamError):
+        f((1, 2), ({"a": "hi"}, {"a": "bye"}, {"a": "oh-no"}))
+
+    with pytest.raises(ParamError):
+        f((3, 4, 5), {"a": ("hi", "bye")})
+
+
+def test_structure_suffix(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[int, " T"], y: PyTree[str, "... T"]):
+        pass
+
+    f(1, "hi")
+    f((3, 4, {"a": 5}), ("a", "b", {"a": "c"}))
+    f(1, ("hi",))
+
+    with pytest.raises(ParamError):
+        f((3, 4), {"a": (1, 2)})
+
+    with pytest.raises(ParamError):
+        f((3, 4, 5), {"a": ("hi", "bye")})
+
+
+def test_structure_compose(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[int, " T"], y: PyTree[int, " S"], z: PyTree[str, "S T"]):
+        pass
+
+    f(1, 2, "hi")
+    f((1, 2), 2, ("a", "b"))
+
+    with pytest.raises(ParamError):
+        f((1, 2), 2, (1, 2))
+
+    f((1, 2), {"a": 3}, {"a": ("hi", "bye")})
+
+    with pytest.raises(ParamError):
+        f((1, 2), {"a": 3}, ({"a": "hi"}, {"a": "bye"}))
+
+    @jaxtyp(typecheck)
+    def g(x: PyTree[int, " T"], y: PyTree[int, " S"], z: PyTree[str, "T S"]):
+        pass
+
+    with pytest.raises(ParamError):
+        g((1, 2), {"a": 3}, {"a": ("hi", "bye")})
+
+    g((1, 2), {"a": 3}, ({"a": "hi"}, {"a": "bye"}))
+
+
+@pytest.mark.parametrize("variadic", (False, True))
+def test_treepath_dependence_function(variadic, jaxtyp, typecheck, getkey):
+    if variadic:
+        jtshape = "*?foo"
+        shape = (2, 3)
+    else:
+        jtshape = "?foo"
+        shape = (4,)
+
+    @jaxtyp(typecheck)
+    def f(
+        x: PyTree[Float[Array, jtshape], " T"], y: PyTree[Float[Array, jtshape], " T"]
+    ):
+        pass
+
+    x1 = jr.normal(getkey(), shape)
+    y1 = jr.normal(getkey(), shape)
+    x2 = jr.normal(getkey(), (5,))
+    y2 = jr.normal(getkey(), (5,))
+    f(x1, y1)
+    f((x1, x2), (y1, y2))
+
+    with pytest.raises(ParamError):
+        f(x1, y2)
+
+    with pytest.raises(ParamError):
+        f((x1, x2), (y2, y1))
+
+
+@pytest.mark.parametrize("variadic", (False, True))
+def test_treepath_dependence_dataclass(variadic, typecheck, getkey):
+    if variadic:
+        jtshape = "*?foo"
+        shape = (2, 3)
+    else:
+        jtshape = "?foo"
+        shape = (4,)
+
+    @jaxtyping.jaxtyped(typechecker=typecheck)
+    class A(eqx.Module):
+        x: PyTree[Float[Array, jtshape], " T"]
+        y: PyTree[Float[Array, jtshape], " T"]
+
+    x1 = jr.normal(getkey(), shape)
+    y1 = jr.normal(getkey(), shape)
+    x2 = jr.normal(getkey(), (5,))
+    y2 = jr.normal(getkey(), (5,))
+    A(x1, y1)
+    A((x1, x2), (y1, y2))
+
+    with pytest.raises(ParamError):
+        A(x1, y2)
+
+    with pytest.raises(ParamError):
+        A((x1, x2), (y2, y1))
+
+
+def test_treepath_dependence_missing_structure_annotation(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[Float[Array, "?foo"], " T"], y: PyTree[Float[Array, "?foo"]]):
+        pass
+
+    x1 = jr.normal(getkey(), (2,))
+    y1 = jr.normal(getkey(), (2,))
+    with pytest.raises(ValueError, match="except when contained with structured"):
+        f(x1, y1)
+
+
+def test_treepath_dependence_multiple_structure_annotation(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[PyTree[Float[Array, "?foo"], " S"], " T"]):
+        pass
+
+    x1 = jr.normal(getkey(), (2,))
+    with pytest.raises(ValueError, match="ambiguous which PyTree"):
+        f(x1)

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_tf_dtype.py

@@ -0,0 +1,13 @@
+# Tensorflow dependency kept in a separate file, so that we can optionally exclude it
+# more easily.
+import tensorflow as tf
+
+from jaxtyping import UInt
+
+
+def test_tf_dtype():
+    x = tf.constant(1, dtype=tf.uint8)
+    y = tf.constant(1, dtype=tf.float32)
+    hint = UInt[tf.Tensor, "..."]
+    assert isinstance(x, hint)
+    assert not isinstance(y, hint)

test/test_threading.py

@@ -31,8 +31,6 @@ class _ErrorableThread(threading.Thread):
             super().run()
         except Exception as e:
             self.exc = e
-        finally:
-            del self._target, self._args, self._kwargs
 
     def join(self, timeout=None):
         super().join(timeout)