Open thejcannon opened 3 months ago
Hey there!
Hmm, I don't think I understand this one. Doing just python -c 'import jaxtyping'
shouldn't load pytest at all. In fact with this not even jaxtyping._pytest_plugin
is loaded -- this entry point is accessed only by pytest itself.
Right, the two things I was trying to show are:
jaxtyping
takes 2.75sjaxtyping
has a pytest hook, pytest loads it, and because the hook points to jaxtyping._pytest_plugin
, which is underneath jaxtyping
, then jaxtyping/__init__.py
is loaded (and he's the slow boy)So, any pytest
invocation is now 2.75s slower.
Ah, hmm. Perhaps it'll be the way that jaxtyping
checks if JAX/Equinox are available. I'm speculating here as jaxtyping
itself is tiny.
If you can, try benchmarking what happens when {jax, equinox}
x {are, aren't}
available to be imported and see how that changes things? Maybe we can optimise that somehow.
(FWIW I frequently use jaxtyping + pytest
and don't see this kind of multisecond lag, but maybe there's some kind of machine-to-machine variation there.)
👋
On my machine:
so, because
jaxtyping
has the pytest entry point, pytest will loadjaxtyping._pytest_plugin
which loadsjaxtyping
which takes 2.75s.Would you consider putting the pytest-related code under a different tree?