mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-08-24 12:11:18 +08:00
Patrick Kidger
91a36aaee4
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).