patrick-kidger/jaxtyping v0.2.25

jaxtyping v0.2.25

on GitHub

This release is primarily a usability release, designed to help ensure the library is being used correctly.

• The error messages from a failed typecheck have been improved, to explicitly highlight more information about which argument was wrong. :)
• If the jaxtyping.jaxtyped(typechecker=...) argument is not passed, then a warning will be displayed. In practice, this will trigger:
  • if using the old double-decorator syntax (@jaxtyped @beartype def foo(...): ...) -- upgrade to the new @jaxtyped(typechecker=beartype) def foo(...): ... syntax and get better error messages! :)
  • If making the easy mistake of writing @jaxtyped(beartype) def foo(...): ... -- in this case it's actually the beartype call that is jaxtype'd, not foo.
• Incorrect use of jaxtyping annotations will now raise an jaxtyping.AnnotationError rather than a mix of RuntimeErrors, NameErrors etc. For example isinstance(x, Float) is not correct (you should write something like Float[Array, "..."]) instead), and this will raise such an AnnotationError.
• Introduced two config flags:
  • JAXTYPING_DISABLE=1 / jaxtyping.config.update("jaxtyping_disable", True): if enabled then all runtime type checking will be skipped.
  • JAXTYPING_REMOVE_TYPECHECKER_STACK=1 / jaxtyping.config.update("jaxtyping_remove_typechecker_stack", True): if enabled then type-checking errors will only show the jaxtyping.TypeCheckError, and won't include any extra stack trace from the underlying type-checker (beartype/typeguard). Some users have found that they preferred the conciseness over the extra information.

Full Changelog: v0.2.24...v0.2.25