[Flax] Add jax flax to env command (#12251)
[Flax] Add jax flax to env command (#12251)

* fix_torch_device_generate_test

* remove @

* add commands for flax/jax
huggingface/transformers,huggingface/transformers,huggingface/transformers,huggingface/transformers