Fix detection of epath

[NeilGirdhar]

Unfortunately, the old detection code doesn't guarantee that epath is installed:

[utM] In [7]: importlib.util.find_spec("etils.epath")
Out[7]: ModuleSpec(name='etils.epath',
loader=<_frozen_importlib_external.SourceFileLoader object at
0x73b8492a7230>,
origin='/home/neil/src/cmm/.venv/lib/python3.12/site-packages/etils/epath/__init__.py',
submodule_search_locations=['/home/neil/src/cmm/.venv/lib/python3.12/site-packages/etils/epath'])

[utM] In [8]: import etils.epath
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent
call last)
Cell In[8], line 1
----> 1 import etils.epath
...
ModuleNotFoundError: No module named 'importlib_resources'

This happened every time I ran jax with a clean environment.

[superbobry]

Thanks, Neil!

[NeilGirdhar]

I replaced epath_installed to fix the build errors, and got rid of the lazy loading since it doesn't do anything anymore. (If the lazy loading is desired, epath may need to be changed so that the old importlib tests work.)

[superbobry]

I think we do want path instead of pathlib here.

[NeilGirdhar]

Why do you want that? Whether you use etils or not, you get a pathlib.Path (essentially) instance here.

The issue is that jax._src.path.Path is now Callable[[...], pathlib.Path], which can't be used as an annotation.

I had to make it a callable rather than type[pathlib.Path] because in Python, initializers don't obey LSP, so type[pathlib.Path]'s __init__ is essentially untyped under some type checkers. (I wrote about this, and it appears that MyPy has recently started switching its stance on this based on comments by their team, but I don't have links to that.)

[superbobry]

Actually, ignore me, this is just an interface so pathlib.Path should be fine. Sorry for the noise and thank you for the additional context.

[NeilGirdhar]

Thank you, always a pleasure interacting with the Jax team!