patrick-kidger / jaxtyping

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/
Other
1.18k stars 62 forks source link

Define the IPython magic lazily #201

Closed superbobry closed 6 months ago

superbobry commented 6 months ago

IPython is quite a chunky package, and importing it unconditionally makes jaxtyping itself relatively slow to import.

patrick-kidger commented 6 months ago

This looks good to me! Thank you for the improvement. (Also CC @knyazer who wrote this originally.)

knyazer commented 6 months ago

Yep, LGTM.

Thanks for the CC, Patrick :)