patrick-kidger/jaxtyping v0.3.0

jaxtyping v0.3.0

on GitHub

Features

• Added support for MLX arrays. (Thanks @gabrieldemarmiesse! #299)
• Added support for int2 and uint2 datatypes (Thanks @mattjj! #307)

(Very mildly) breaking changes

• Previously we arranged it such that jaxtyping annotations were a subclass of their array type, e.g. Float[Array, "foo"] was a subclass of Array. This ran into various issues with un-subclass-able types, metaclass conflicts, etc. Correspondingly this subclassing has been removed.
• @jaxtyped-decorated dataclasses now check the arguments to their __init__ methods, rather than checking their attributes after initialisation. (This better handles some edge-cases with dataclass extensions where fields are allowed to have converters for their type.)

Misc

• More efficient @jaxtyped decoration (Thanks @hawkinsp! #311)
• Better pretty-printed representations for PyTree.
• Prettier documentation :)

New Contributors

• @mattjj made their first contribution in #307
• @hawkinsp made their first contribution in #311

Full Changelog: v0.2.38...v0.3.0