AttributeError: module 'jax.random' has no attribute 'KeyArray' while fine tuning. #221

samyakai · 2022-04-20T12:08:59Z

I am following your rep to fine tune GPT-J on TPU. When I run "python3 device_train.py --config=YOUR_CONFIG.json --tune-model-path=gs://YOUR-BUCKET/step_383500/" with my bucket name and the config file I have created, I get an error as "AttributeError: module 'jax.random' has no attribute 'KeyArray'". These are some of the specs:

OS: Ubuntu 20.04
jax version = 0.2.12
TPU : V3-8
Zone : us-central1-b

The error is caused by line 7 in the device_train.py where optax is being imported:- "import optax".

This is the error stack:

WARNING: Logging before InitGoogle() is written to STDERR
I0420 11:47:44.856002 10240 tpu_initializer_helper.cc:94] libtpu.so already in use by another process. Run "$ sudo lsof -w /dev/accel0" to figure out which process is using the TPU. Not attempting to load libtpu.so in this process.
Traceback (most recent call last):
File "device_train.py", line 7, in
import optax
File "/usr/local/lib/python3.8/dist-packages/optax/init.py", line 17, in
from optax._src.alias import adabelief
File "/usr/local/lib/python3.8/dist-packages/optax/_src/alias.py", line 21, in
from optax._src import base
File "/usr/local/lib/python3.8/dist-packages/optax/_src/base.py", line 18, in
import chex
File "/usr/local/lib/python3.8/dist-packages/chex/init.py", line 17, in
from chex._src.asserts import assert_axis_dimension
File "/usr/local/lib/python3.8/dist-packages/chex/_src/asserts.py", line 26, in
from chex._src import asserts_internal as _ai
File "/usr/local/lib/python3.8/dist-packages/chex/_src/asserts_internal.py", line 32, in
from chex._src import pytypes
File "/usr/local/lib/python3.8/dist-packages/chex/_src/pytypes.py", line 36, in
PRNGKey = jax.random.KeyArray
AttributeError: module 'jax.random' has no attribute 'KeyArray'

Any help is appreciated!

mosmos6 · 2022-04-20T12:10:58Z

I've just encountered exactly the same error and I was about to open an issue about this.

KD1903 · 2022-04-20T12:13:09Z

WARNING: Logging before InitGoogle() is written to STDERR
I0420 11:47:44.856002 10240 tpu_initializer_helper.cc:94] libtpu.so already in use by another process. Run "$ sudo lsof -w /dev/accel0" to figure out which process is using the TPU. Not attempting to load libtpu.so in this process.
Traceback (most recent call last):
File "device_train.py", line 7, in
import optax
File "/usr/local/lib/python3.8/dist-packages/optax/init.py", line 17, in
from optax._src.alias import adabelief
File "/usr/local/lib/python3.8/dist-packages/optax/_src/alias.py", line 21, in
from optax._src import base
File "/usr/local/lib/python3.8/dist-packages/optax/_src/base.py", line 18, in
import chex
File "/usr/local/lib/python3.8/dist-packages/chex/init.py", line 17, in
from chex._src.asserts import assert_axis_dimension
File "/usr/local/lib/python3.8/dist-packages/chex/_src/asserts.py", line 26, in
from chex._src import asserts_internal as _ai
File "/usr/local/lib/python3.8/dist-packages/chex/_src/asserts_internal.py", line 32, in
from chex._src import pytypes
File "/usr/local/lib/python3.8/dist-packages/chex/_src/pytypes.py", line 36, in
PRNGKey = jax.random.KeyArray
AttributeError: module 'jax.random' has no attribute 'KeyArray'

I am facing the same error. Kindly solve this!

jagruti-samyak · 2022-04-20T12:18:54Z

The same error is facing while "import optax".

WARNING: Logging before InitGoogle() is written to STDERR
I0420 11:47:44.856002 10240 tpu_initializer_helper.cc:94] libtpu.so already in use by another process. Run "$ sudo lsof -w /dev/accel0" to figure out which process is using the TPU. Not attempting to load libtpu.so in this process.
Traceback (most recent call last):
File "device_train.py", line 7, in
import optax
File "/usr/local/lib/python3.8/dist-packages/optax/init.py", line 17, in
from optax._src.alias import adabelief
File "/usr/local/lib/python3.8/dist-packages/optax/_src/alias.py", line 21, in
from optax._src import base
File "/usr/local/lib/python3.8/dist-packages/optax/_src/base.py", line 18, in
import chex
File "/usr/local/lib/python3.8/dist-packages/chex/init.py", line 17, in
from chex._src.asserts import assert_axis_dimension
File "/usr/local/lib/python3.8/dist-packages/chex/_src/asserts.py", line 26, in
from chex._src import asserts_internal as _ai
File "/usr/local/lib/python3.8/dist-packages/chex/_src/asserts_internal.py", line 32, in
from chex._src import pytypes
File "/usr/local/lib/python3.8/dist-packages/chex/_src/pytypes.py", line 36, in
PRNGKey = jax.random.KeyArray
AttributeError: module 'jax.random' has no attribute 'KeyArray'

abdelatifsd · 2022-04-20T12:34:48Z

Exact same issue here!

hxiaoyang · 2022-04-20T14:30:59Z

same issue!

vfbd · 2022-04-20T17:00:24Z

Chex 0.1.3 doesn't support JAX 0.2.12. You need to downgrade to Chex 0.1.2:

pip3 install chex==0.1.2

mosmos6 · 2022-04-21T06:50:06Z

@vfbd It worked for me to infer the model. but apparently not for finetuning.

samyakai · 2022-04-21T07:41:23Z

@mosmos6 @vfbd Now it is giving me this error: "AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'".
@mosmos6 How did it work for you? Are you training on TPU v3-8?

mosmos6 · 2022-04-21T08:15:37Z

@samyakai Now I noticed you encountered this error on fine tune. I did on inference but the same error. Sorry for the confusion. I modified my previous comment. The issue hasn't been resolved for finetuning.

samyakai · 2022-04-21T10:22:30Z

As suggested by @vfbd if I downgrade chex to 0.1.2 I encounter " "AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'".". To overcome this https://github.com/google/brax/issues/187 suggests upgrading to latest version. if I do that I again encounter the error "AttributeError: module 'jax.random' has no attribute 'KeyArray'" .

jagruti-samyak · 2022-04-21T10:30:15Z

I am following your rep to fine tune GPT-J on TPU. When I run "python3 device_train.py --config=YOUR_CONFIG.json --tune-model-path=gs://YOUR-BUCKET/step_383500/" with my bucket name and the config file I have created, I get an error as "AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice''". These are some of the specs:

OS: Ubuntu 20.04
jax version = 0.2.12
chex version == 0.1.2
TPU : V3-8
Zone : us-central1-b

The error is caused by line 7 in the device_train.py where optax is being imported:- "import optax".

This is the error stack:

WARNING: Logging before InitGoogle() is written to STDERR
I0421 10:06:19.047791 8679 tpu_initializer_helper.cc:94] libtpu.so already in use by another process. Run "$ sudo lsof -w /dev/accel0" to figure out which pr ocess is using the TPU. Not attempting to load libtpu.so in this process.
Traceback (most recent call last):
File "device_train.py", line 7, in
import optax
File "/usr/local/lib/python3.8/dist-packages/optax/init.py", line 17, in < module>
from optax import experimental
File "/usr/local/lib/python3.8/dist-packages/optax/experimental/init.py", line 20, in
from optax._src.experimental.complex_valued import split_real_and_imaginary
File "/usr/local/lib/python3.8/dist-packages/optax/_src/experimental/complex_v alued.py", line 32, in
import chex
File "/usr/local/lib/python3.8/dist-packages/chex/init.py", line 17, in
from chex._src.asserts import assert_axis_dimension
File "/usr/local/lib/python3.8/dist-packages/chex/_src/asserts.py", line 26, i n
from chex._src import asserts_internal as _ai
File "/usr/local/lib/python3.8/dist-packages/chex/_src/asserts_internal.py", l ine 32, in
from chex._src import pytypes
File "/usr/local/lib/python3.8/dist-packages/chex/_src/pytypes.py", line 40, i n
CpuDevice = jax.lib.xla_extension.CpuDevice
AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'

Please help us to resolve it asap..
Thank you

vfbd · 2022-04-22T05:16:59Z

Which version of jaxlib (not jax) do you have? Maybe try again with jaxlib==0.1.68

samyakai · 2022-04-22T09:34:22Z

@vfbd These are the library versions which solve the error.
jax==0.2.16
jaxlib==0.1.68
optax==0.1.2
chex==0.1.2

kufton · 2023-01-21T05:17:33Z

I just paid for colab pro to play around with this and found the same issues described here. I added !pip install for the lib versions mentioned and then I got this error:

`---------------------------------------------------------------------------
ImportError Traceback (most recent call last)
in
4 from jax.experimental import maps
5 import numpy as np
----> 6 import optax
7 import transformers
8

6 frames
/usr/local/lib/python3.8/dist-packages/jax/_src/api.py in
42 from . import dtypes
43 from ..core import eval_jaxpr
---> 44 from ..api_util import (flatten_fun, apply_flat_fun, flatten_fun_nokwargs,
45 flatten_fun_nokwargs2, argnums_partial,
46 argnums_partial_except, flatten_axes, donation_vector,

ImportError: cannot import name '_ensure_str_tuple' from 'jax.api_util' (/usr/local/lib/python3.8/dist-packages/jax/api_util.py)

NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
"Open Examples" button below.
---------------------------------------------------------------------------`

I appreciate that this is like, alpha, so while I'll go play with GTP3, thank you for your work.

mystiverv · 2024-04-05T00:51:43Z

from jax_md import rigid_body
File "C:\Users....\env\Lib\site-packages\jax_md\rigid_body.py", line 76, in
KeyArray = random.KeyArray

module 'jax.random' has no attribute 'KeyArray'

I get this error when trying to run import jax_md

samyakai changed the title ~~AttributeError: module 'jax.random' has no attribute 'KeyArray'~~ AttributeError: module 'jax.random' has no attribute 'KeyArray' while fine tuning. Apr 21, 2022

samyakai closed this as completed Apr 23, 2022

lightyrs mentioned this issue Apr 28, 2022

The latest update to PIP breaks installation #210

Open

vfbd mentioned this issue Jul 8, 2022

Google Colab Error: optax is throwing an attribute error. #230

Open

kevinrue mentioned this issue Feb 7, 2024

import cell2location -> AttributeError: module 'jax.random' has no attribute 'KeyArray' BayraktarLab/cell2location#347

Open

2 tasks

heyitsguay mentioned this issue Feb 14, 2024

Requirements version error LargeWorldModel/LWM#3

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

AttributeError: module 'jax.random' has no attribute 'KeyArray' while fine tuning. #221

AttributeError: module 'jax.random' has no attribute 'KeyArray' while fine tuning. #221

samyakai commented Apr 20, 2022

mosmos6 commented Apr 20, 2022

KD1903 commented Apr 20, 2022

jagruti-samyak commented Apr 20, 2022

abdelatifsd commented Apr 20, 2022

hxiaoyang commented Apr 20, 2022

vfbd commented Apr 20, 2022

mosmos6 commented Apr 21, 2022 •

edited

Loading

samyakai commented Apr 21, 2022 •

edited

Loading

mosmos6 commented Apr 21, 2022 •

edited

Loading

samyakai commented Apr 21, 2022

jagruti-samyak commented Apr 21, 2022

vfbd commented Apr 22, 2022 •

edited

Loading

samyakai commented Apr 22, 2022 •

edited

Loading

kufton commented Jan 21, 2023

mystiverv commented Apr 5, 2024

AttributeError: module 'jax.random' has no attribute 'KeyArray' while fine tuning. #221

AttributeError: module 'jax.random' has no attribute 'KeyArray' while fine tuning. #221

Comments

samyakai commented Apr 20, 2022

mosmos6 commented Apr 20, 2022

KD1903 commented Apr 20, 2022

jagruti-samyak commented Apr 20, 2022

abdelatifsd commented Apr 20, 2022

hxiaoyang commented Apr 20, 2022

vfbd commented Apr 20, 2022

mosmos6 commented Apr 21, 2022 • edited Loading

samyakai commented Apr 21, 2022 • edited Loading

mosmos6 commented Apr 21, 2022 • edited Loading

samyakai commented Apr 21, 2022

jagruti-samyak commented Apr 21, 2022

vfbd commented Apr 22, 2022 • edited Loading

samyakai commented Apr 22, 2022 • edited Loading

kufton commented Jan 21, 2023

mystiverv commented Apr 5, 2024

module 'jax.random' has no attribute 'KeyArray'

mosmos6 commented Apr 21, 2022 •

edited

Loading

samyakai commented Apr 21, 2022 •

edited

Loading

mosmos6 commented Apr 21, 2022 •

edited

Loading

vfbd commented Apr 22, 2022 •

edited

Loading

samyakai commented Apr 22, 2022 •

edited

Loading