jax hybrip_clip error ( numpy  is not a valid JAX type )
Author: acul3
Pull request: no

i tried to run  script([run_hybrid_clip.py](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects/hybrid_clip))

after few step i got error

```
it/s/home/acul/env/lib/python3.8/site-packages/transformers/tokenization_utils_base.py:699: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.
  tensor = as_tensor(value)
Epoch ... (1/40):   0%|                                                | 0/40 [02:43<?, ?it/s]
Traceback (most recent call last):                           | 39/462 [02:43<06:48,  1.03it/s]
  File "run_hybrid_clip.py", line 556, in <module>
    main()
  File "run_hybrid_clip.py", line 502, in main
    state, train_metric = p_train_step(state, batch)
  File "/home/acul/env/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/acul/env/lib/python3.8/site-packages/jax/_src/api.py", line 1623, in f_pmapped
    for arg in args: _check_arg(arg)
  File "/home/acul/env/lib/python3.8/site-packages/jax/_src/api.py", line 2296, in _check_arg
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")
jax._src.traceback_util.UnfilteredStackTrace: TypeError: Argument '[[list([1, 1, 1, 1, 
.
.
.
.
  list([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])]]' of type <class 'numpy.ndarray'> is not a valid JAX type.

```
i use [bert](https://huggingface.co/indobenchmark/indobert-base-p1) and clip as model vision and text


command that i use 
```
python run_hybrid_clip.py \
    --output_dir ${MODEL_DIR} \
    --text_model_name_or_path="indobenchmark/indobert-base-p1" \
    --vision_model_name_or_path="openai/clip-vit-base-patch32" \
    --tokenizer_name="indobenchmark/indobert-base-p1" \
    --train_file="coco_dataset/train_dataset.json" \
    --validation_file="coco_dataset/validation_dataset.json" \
    --do_train --do_eval \
    --num_train_epochs="40" --max_seq_length 96 \
    --per_device_train_batch_size="64" \
    --per_device_eval_batch_size="64" \
    --learning_rate="5e-5" --warmup_steps="0" --weight_decay 0.1 \
    --overwrite_output_dir \
    --preprocessing_num_workers 32 \
    --push_to_hub
```

i also use same dataset as the example 
pinging @patil-suraj 

any idea how to fix ?

Comments:
Looking into it.

I just tried it, and couldn't reproduce the issue. Could you please verify again?

Also, note that in COCO the captions are in English so this won't work for  Indonesian. You should use a dataset that has captions in Indonesian. The [WIT](https://github.com/google-research-datasets/wit) dataset might help.

fyi i translated coco dataset to indonesian using [EasyNmt](https://github.com/UKPLab/EasyNMT)

it seems the problem is between my translate coco dataset or the jax version in my tpu vm
because when i  run the code in my local machine(with "jax[cuda111]" ) using the same dataset, the training run smoothly

will try to explore more to confirm  issue!

this problem occurs when a sequence longer than `max_seq_length` appears, since the defualt tokenization function does not truncate the sequences. 
Try increasing the `max_seq_len` (preferably to a multiple of 64, in order to improve TPU performance) or adding the truncation in the tokenization function