AttributeError: module 'jax.numpy' has no attribute 'DeviceArray' in colab
Author: yundaehyuck
Pull request: no

### System Info

transformers - 4.31.0
python - 3.10.12
colab

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

i run this code in colab, but an issue arises.

3 days ago, i run this code normally.

why is this error occuring?

!pip install transformers

from transformers import *

/usr/local/lib/python3.10/dist-packages/transformers/generation_utils.py:24: FutureWarning: Importing `GenerationMixin` from `src/transformers/generation_utils.py` is deprecated and will be removed in Transformers v5. Import as `from transformers import GenerationMixin` instead.
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/transformers/generation_tf_utils.py:24: FutureWarning: Importing `TFGenerationMixin` from `src/transformers/generation_tf_utils.py` is deprecated and will be removed in Transformers v5. Import as `from transformers import TFGenerationMixin` instead.
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/transformers/generation_flax_utils.py:24: FutureWarning: Importing `FlaxGenerationMixin` from `src/transformers/generation_flax_utils.py` is deprecated and will be removed in Transformers v5. Import as `from transformers import FlaxGenerationMixin` instead.
  warnings.warn(
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[/usr/local/lib/python3.10/dist-packages/transformers/utils/import_utils.py](https://localhost:8080/#) in _get_module(self, module_name)
   1098         try:
-> 1099             return importlib.import_module("." + module_name, self.__name__)
   1100         except Exception as e:

15 frames
AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
[/usr/local/lib/python3.10/dist-packages/transformers/utils/import_utils.py](https://localhost:8080/#) in _get_module(self, module_name)
   1099             return importlib.import_module("." + module_name, self.__name__)
   1100         except Exception as e:
-> 1101             raise RuntimeError(
   1102                 f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"
   1103                 f" traceback):\n{e}"

RuntimeError: Failed to import transformers.models.bart.modeling_flax_bart because of the following error (look up to see its traceback):
module 'jax.numpy' has no attribute 'DeviceArray'

### Expected behavior

no

Comments:
cc @sanchit-gandhi 

I am getting the same error (not in colab).

```python
from transformers.models.longt5 import modeling_flax_longt5
```

Returns:
```
RuntimeError: Failed to import transformers.models.longt5.modeling_flax_longt5 because of the following error (look up to see its traceback):
module 'jax.numpy' has no attribute 'DeviceArray'
```

Versions:
```
transformers==4.27.0
jax==0.4.14
torch==1.12.1
```

Just saw this PR: https://github.com/huggingface/transformers/pull/24875 Will try the main branch.

@HarshTrivedi @sgugger 

thank you for comment

I check the PR #24875. 

!pip install jax==0.4.13
!pip install jaxlib==0.4.13

When these codes are executed in colab, transformers works normally.

Until the new version comes out, I will use it like this.

thanks.