Update `Array` type annotation to use the new jax typing module.

LouisDesdoigts / dLux

Differentiable optical models as parameterised neural networks in Jax using Zodiax

https://louisdesdoigts.github.io/dLux/

BSD 3-Clause "New" or "Revised" License

43 stars 6 forks source link

Closed LouisDesdoigts closed 1 year ago

LouisDesdoigts commented 1 year ago

Patrick Kidger's JaxTyping module was merged into jax! We should update to use this internally https://jax.readthedocs.io/en/latest/jax.typing.html#module-jax.typing. This should be a simple import change.

LouisDesdoigts commented 1 year ago

Done in 0.13