* Add a test for generators
* Remove output annotations from decorators
Also guarded torch imports for better compatibility with
requirements.txt
* Add flag to the main meta class to skip the typecheck
* Return to the old solution
* Make async tests work
* Minor adjustments/fixing typos
* Correct Python path for new tests
* Remove some jax-dependent code
* Implement equality for MetaArrays
* Make all Dim variations frozen dataclasses
* Shorten AbstractArray methods
* Final touches
* Removing get_origin use
* Update tests with @jaxtyp
Recall that jaxtyping will currently generate rich error messages in precisely one scenario: about the arguments and return types when doing:
```python
@jaxtyped(typechecker=beartype)
def foo(...): ...
```
With this commit we add support for beartype 0.17.0's pseudo-standard `__instancecheck_str__`, which means the following:
1. For those using beartype decorators, the following will *also* generate an informative error message, and moreover it will state exactly why (shape mismatch, dtype mismatch etc):
```python
@jaxtyped(typechecker=None)
@beartype
def foo(...): ...
```
(In practice we probably won't recommend the above combination in the docs just to keep things simple.)
2. For those using the beartype import hook together with the jaxtyping import hook, we can probably also check `assert isinstance(x, Float[Array, "foo"])` statements with rich error messages. (#153) We'll need to test + document that though. (@jeezrick interested?)
3. For those using plain `assert isinstance(...)` statements without beartype (#167, tagging @reinerp), then they can *also* get rich error messages by doing
```python
tt = Float[Array, "foo"]
assert isinstance(x, tt), tt.__instancecheck_str__(x) + "\n" + print_bindings()
```
which is still a bit long-winded right now but is a step in the right direction.
(CC @leycec for interest.)
Hello!
This is a small PR to fix typos in docstrings.
Maybe, I would suggest adding an import for `dataclass` in the example (otherwise it will not run), and maybe indicate that it works with other dataclasses decorators, like the `dataclass` decorator from chex.
These flags are `JAXTYPING_DISABLE` and `JAXTYPING_REMOVE_TYPECHECKER_STACK`.
In addition, have now added warnings when using old-style double-decorator syntax, which also serves to guard against the easy mistake of
```python
@jaxtyped(typechecker)
def foo(...)
```
which actually decorates the `typechecker`, not `foo`.
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?)
Previously, something like this would not raise an error, as
variadic+broadcast dimensions were stored in a separate namespace to
variadic+nonbroadcast dimensions:
```python
def f(x: Float[Array, "*foo"], y: Float[Array, "#*foo"]):
pass
a, b = ...
assert a.shape == (3, 4)
assert b.shape == (5,)
f(a, b)
```
Michael J Clark