* 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
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)
```
Previously, using the import hook with dataclasses resulted in the `__init__` method of the dataclass being checked.
This was undesirable when using `eqx.field(converter=...)`, as the annotation didn't necessarily reflect the argument type.
A typical example was
```python
class Foo(eqx.Module):
x: jax.Array = eqx.field(converter=jnp.ndarray)
Foo(1) # 1 is not an array! But this code is valid.
```
After this change, we instead monkey-patch our checks to happen at the end of the `__init__` of the dataclass -- after conversion has run.
Note that this requires https://github.com/patrick-kidger/equinox/pull/524. Otherwise, Equinox does conversion too late (in `_ModuleMeta.__call__`, after `__init__` has been run).
* Various improvements.
- Added support for functions in symbolic dimensions, e.g. "min(foo,bar)", which were previously disallowed due to the presence of a comma. (#51)
- Added support for adding ignored names to dimensions, e.g. "cols=4". (#76)
* Now works with Python 3.10 A | B union types.
This required quite a lot of refactoring! JAX supports virtual subclass registration (its metaclass is ABCMeta) but NumPy does not, so we have to actually subclass `np.ndarray`.
Simple stuff like __base__ hacking fails due to deallocator conflicts.
- Float32 introduced to replace f32 (etc.)
- Old f32 aliases were previously around for backward-compatibility, but
the code is pretty hideous, and all stakeholders are now on board
through RFC #13.
- Overall feedback was that Int{Sign,Unsign} was too long and not enough
like numpy. Changed to Int, UInt, and Integer.
Roman Knyazhitskiy