Misleading ImportError when using JAX tensors without Flax installed
Author: yixiaoer
Pull request: no
Labels: name: WIP

name: jax

### System Info

- `transformers` version: 4.37.2
- Platform: Linux-5.19.0-1027-gcp-x86_64-with-glibc2.35
- Python version: 3.11.4
- Huggingface_hub version: 0.20.1
- Safetensors version: 0.4.1
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.3.0.dev20231228+cpu (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed  # But I have installed
- JaxLib version: not installed  # But I have installed
- Using GPU in script?: N/A
- Using distributed or parallel set-up in script?: N/A

### Who can help?

@sanchit-gandhi @ArthurZucker

### 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

While attempting to convert tokenizer outputs to JAX tensors using the following code:
```python
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('mistralai/Mistral-7B-v0.1')
sentences = ['hello world']
inputs = tokenizer(sentences, padding=True, return_tensors='jax')
```
I received the following ImportError:
```
ImportError: Unable to convert output to JAX tensors format, JAX is not installed.
```

However, JAX is indeed installed in my environment.

### Expected behavior

Upon further investigation, it seems the error arises because the library checks for Flax's availability (`is_flax_available()`) rather than JAX's direct presence. Here is the snippet from the [source code](https://github.com/huggingface/transformers/blob/58e3d23e97078f361a533b9ec4a6a2de674ea52a/src/transformers/tokenization_utils_base.py#L723-L724) that led to this conclusion:
```python
if not is_flax_available():
    raise ImportError("Unable to convert output to JAX tensors format, JAX is not installed.")
```
This can be somewhat misleading, as the error message suggests a lack of JAX installation, while the actual requirement is for Flax. Not all JAX users utilize Flax, and this might cause confusion.

Would it be possible to update the error message to more accurately reflect the requirement for Flax when attempting to use JAX tensor formats? Such a clarification would greatly assist users in diagnosing setup issues.

Comments:
Would you like to open a PR to have a nicer failure?

Sure! I will do it.

I've read through the relevant source code and realised that there isn't a trivial fix for the issue.

* Use Flax with ```USE_JAX = os.environ.get("USE_FLAX", "AUTO").upper()``` to show the availability of JAX.

* In some parts of the code, `is_flax_available()` was called where only JAX was required, for example:
```python
if is_torch_available():
    import torch
if is_tf_available():
    import tensorflow as tf
if is_flax_available():
    import jax.numpy as jnp  # noqa: F401
from .pipelines.conversational import Conversation
```
* In some parts of the code, `is_flax_available()` was used to detect the existence of Flax, but "JAX is not installed" was raised as the error message.

It appears that the current repository does not make a clear distinction between JAX and Flax. I suggest that the repository distinguish the JAX and Flax in this repo to reduce confusion and improve debugging experiences for users.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

This issue still need to be addressed.

@ArthurZucker Is there an official maintainer working on it?

Nope, as I mentioned before, not super high priority for us given the low usage of jax in transformers ! Do you want to tackle this?