[BUG] fail to run sglang_jax_lora_test.py with flax==0.12.5

[aolemila]

Expected Behavior

Pass the CI.

Actual Behavior

Failed.

Steps to Reproduce the Problem

1. Run python3 tests/generate/sglang_jax_lora_test.py

Environment

• OS: [e.g., Ubuntu, etc.] tpu-v6e-1
• Project Version: [e.g., 0.0.1]

Checklist

• I have searched the existing issues for a similar bug report.
• I have provided all the required information in the "Environment" section.
• I have provided a minimal, reproducible example.

Would you like to help us fix it?

Yes. Currently I fix the flax to 0.12.4. I will fix it later. Track it later with #1207.