2

Could anyone please help me fix the following error when going through "/usr/local/lib/python3.7/dist-packages/haiku/_src/transform.py in check_not_jax_transformed(f)"? Thanks a lot.

"AttributeError: module 'jaxlib.xla_extension' has no attribute 'PmapFunction'"
Laura K.
  • 51
  • 2
  • 7

1 Answers1

3

jaxlib.xla_extension.PmapFunction was added in jaxlib version 0.72; it sounds like you have an older jaxlib version installed. You should update it using:

pip install -U jaxlib

Note: if you're using GPU/TPU, you should instead use the appropriate accelerator-specific installation command found at https://github.com/google/jax#installation.

If this does not work, please check your Python version. jaxlib began requiring Python 3.7 or newer in version 0.1.70, so if you are using Python 3.6, you will need to upgrade Python before you can upgrade to a more recent jaxlib.

It appears the problematic line was added to the haiku package a few hours before you posted the question: https://github.com/deepmind/dm-haiku/commit/e6a13af352a8b46d355ac1b7131b64c615cfcf57 Another option if you don't want to update jaxlib would be to install a stable version of dm-haiku rather than using the development version:

pip install dm-haiku==0.0.5
jakevdp
  • 77,104
  • 11
  • 125
  • 160
  • Thank you for your answer. But as I tried "pip install -U jaxlib", the same error occurs. So how to fix it? Thanks. – Laura K. Nov 12 '21 at 09:45
  • See edit - you need to be on Python 3.7 or newer. – jakevdp Nov 12 '21 at 12:45
  • I used Python 3.7 and have been using Python 3.7, but the error still occurs. So how to fix it? Thanks. – Laura K. Nov 12 '21 at 15:05
  • See the rest of the comment: you'll either need to install a newer version of jaxlib, or install an older version of dm-haiku. The pip installation commands in my answer typically will work for this, but if you have a special Python environment setup, they may not. If the `pip` commands do not work, you need to look into how packages are installed on your system and figure out how to install the correct versions. You could start by double-checking what versions of the packages you have installed to see if it's consistent with my suggestions. – jakevdp Nov 12 '21 at 15:25