[Bug] elegy does not work with latest haiku version

[murphyk]

Describe the bug
When I type 'import elegy' I get this error

 File "/home/kpmurphy/mambaforge/lib/python3.10/site-packages/elegy/generalized_module/haiku_module.py", line 4, in <module>
    from haiku._src.base import current_bundle_name

Minimal code to reproduce

import elegy

Expected behavior
A clear and concise description of what you expected to happen.

Library Info
Please provide os info and elegy version.

>> 
>>> jax.__version__
'0.2.28'
>>> haiku.__version__
'0.0.9.dev'
>>> elegy.__version__. #  elegy-0.5.0-py3-none-any.whl 
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
NameError: name 'elegy' is not defined
>>> 

Screenshots

Additional context
Add any other context about the problem here.

[cgarciae]

Hey @murphyk, can you install the latest stable?

pip install elegy==0.8.6

[murphyk]

I get this error

pip install elegy==0.8.6
ERROR: Ignored the following versions that require a different python version: 0.6.0 Requires-Python >=3.6.1,<3.9; 0.7.0 Requires-Python >=3.6.1,<3.9; 0.7.1 Requires-Python >=3.6.1,<3.9; 0.7.2 Requires-Python >=3.6.1,<3.9; 0.7.4 Requires-Python >=3.6.1,<3.9; 0.8.0 Requires-Python >=3.7,<3.10; 0.8.1 Requires-Python >=3.7,<3.10; 0.8.2 Requires-Python >=3.7,<3.10; 0.8.3 Requires-Python >=3.7,<3.10; 0.8.4 Requires-Python >=3.7,<3.10; 0.8.5 Requires-Python >=3.7,<3.10; 0.8.6 Requires-Python >=3.7,<3.10; 0.9.0.dev0 Requires-Python >=3.7,<3.10; 0.9.0.dev1 Requires-Python >=3.7,<3.10; 0.9.0.dev2 Requires-Python >=3.7,<3.10
ERROR: Could not find a version that satisfies the requirement elegy==0.8.6 (from versions: 0.1.0, 0.1.1b0, 0.1.1b1, 0.1.1b2, 0.1.1b3, 0.1.1b4, 0.1.1, 0.1.2, 0.1.3, 0.1.4, 0.1.5, 0.2.0, 0.2.1, 0.2.2, 0.3.0, 0.4.0, 0.4.1, 0.5.0)
ERROR: No matching distribution found for elegy==0.8.6

[murphyk]

When I don't pin the version number I get very old versions of everything

pip install elegy
Collecting elegy
  Using cached elegy-0.5.0-py3-none-any.whl (206 kB)
Collecting jaxlib<0.2.0,>=0.1.59
  Using cached jaxlib-0.1.76-cp310-none-manylinux2010_x86_64.whl (65.2 MB)
Requirement already satisfied: toolz<0.11.0,>=0.10.0 in ./mambaforge/lib/python3.10/site-packages (from elegy) (0.10.0)
Collecting dm-haiku<0.0.3,>=0.0.2
  Using cached dm_haiku-0.0.2-py3-none-any.whl (241 kB)
Requirement already satisfied: numpy<2.0.0,>=1.0.0 in ./mambaforge/lib/python3.10/site-packages (from elegy) (1.23.3)
Collecting jax<0.3.0,>=0.2.9
  Using cached jax-0.2.28-py3-none-any.whl
Requirement already satisfied: pyyaml<6.0.0,>=5.3.1 in ./mambaforge/lib/python3.10/site-packages (from elegy) (5.4.1)
Collecting optax<0.0.2,>=0.0.1
  Using cached optax-0.0.1-py3-none-any.whl (49 kB)
Requirement already satisfied: cloudpickle<2.0.0,>=1.5.0 in ./mambaforge/lib/python3.10/site-packages (from elegy) (1.6.0)
Requirement already satisfied: tabulate<0.9.0,>=0.8.7 in ./mambaforge/lib/python3.10/site-packages (from elegy) (0.8.10)
Requirement already satisfied: tensorboardx<3.0,>=2.1 in ./mambaforge/lib/python3.10/site-packages (from elegy) (2.5.1)
Requirement already satisfied: absl-py>=0.7.1 in ./mambaforge/lib/python3.10/site-packages (from dm-haiku<0.0.3,>=0.0.2->elegy) (1.2.0)
Requirement already satisfied: scipy>=1.2.1 in ./mambaforge/lib/python3.10/site-packages (from jax<0.3.0,>=0.2.9->elegy) (1.9.1)
Requirement already satisfied: opt-einsum in ./mambaforge/lib/python3.10/site-packages (from jax<0.3.0,>=0.2.9->elegy) (3.3.0)
Requirement already satisfied: typing-extensions in ./mambaforge/lib/python3.10/site-packages (from jax<0.3.0,>=0.2.9->elegy) (4.3.0)
Requirement already satisfied: flatbuffers<3.0,>=1.12 in ./mambaforge/lib/python3.10/site-packages (from jaxlib<0.2.0,>=0.1.59->elegy) (2.0.7)
Requirement already satisfied: chex>=0.0.2 in ./mambaforge/lib/python3.10/site-packages (from optax<0.0.2,>=0.0.1->elegy) (0.1.5)
Requirement already satisfied: protobuf<=3.20.1,>=3.8.0 in ./mambaforge/lib/python3.10/site-packages (from tensorboardx<3.0,>=2.1->elegy) (3.19.5)
Requirement already satisfied: dm-tree>=0.1.5 in ./mambaforge/lib/python3.10/site-packages (from chex>=0.0.2->optax<0.0.2,>=0.0.1->elegy) (0.1.7)
Installing collected packages: dm-haiku, jaxlib, jax, optax, elegy
  Attempting uninstall: dm-haiku
    Found existing installation: dm-haiku 0.0.9.dev0
    Uninstalling dm-haiku-0.0.9.dev0:
      Successfully uninstalled dm-haiku-0.0.9.dev0
  Attempting uninstall: jaxlib
    Found existing installation: jaxlib 0.3.20
    Uninstalling jaxlib-0.3.20:
      Successfully uninstalled jaxlib-0.3.20
  Attempting uninstall: jax
    Found existing installation: jax 0.3.21
    Uninstalling jax-0.3.21:
      Successfully uninstalled jax-0.3.21
  Attempting uninstall: optax
    Found existing installation: optax 0.1.3
    Uninstalling optax-0.1.3:
      Successfully uninstalled optax-0.1.3
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
flax 0.6.1 requires jax>=0.3.16, but you have jax 0.2.28 which is incompatible.
Successfully installed dm-haiku-0.0.2 elegy-0.5.0 jax-0.2.28 jaxlib-0.1.76 optax-0.0.1
(base) kpmurphy@t1v-n-6c70f275-w-0:~$ 

[murphyk]

Atfer installing this very old version it does not run , presumaly because it relies on an old tpu driver.
(I am running on a TPU VM V3-8)

base) kpmurphy@t1v-n-6c70f275-w-0:~$ python
Python 3.10.6 | packaged by conda-forge | (main, Aug 22 2022, 20:35:26) [GCC 10.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import elegy as eg
2022-10-08 15:39:57.200035: F external/org_tensorflow/tensorflow/core/tpu/tpu_library_init_fns.inc:101] TpuEmbeddingEngine_ConfigureCommunication not available in this library.
Aborted (core dumped)

[ScottAlexanderCameron]

@murphyk, current version of elegy (see pyproject.toml) requires a python version <3.10. Try something like python 3.9, then you should be able to install elegy 0.8.6