patrick-kidger/jaxtyping v0.2.31

jaxtyping v0.2.31

on GitHub

• Now duck-type on array shapes and dtypes, so you can use jaxtyping for your custom arraylike objects:

   class FooDtype(jaxtyping.AbstractDtype):
   	dtypes = ["foo"]
  
   class MyArray:
   	@property
   	def dtype(self):
   		return "foo"
  
   	@property
   	def shape(self):
   		return (3, 1, 4)
  
   def f(x: FooDtype[MyArray, "3 1 4"]): ...

• Improved compatibility when typeguard warns that you're typechecking a function without annotations: it will no longer mention the jaxtyping-internal check_params function and will instead mention the name of the function that is missing annotations.

• Improved the error message when typechecking fails, to state the full some_module.SomeClass.some_method rather than just some_method.

• Fixed a JAX deprecation warning for jax.tree_map. (Thanks @groszewn!)

New Contributors

• @groszewn made their first contribution in #218

Full Changelog: v0.2.30...v0.2.31