Jax jitのトレースはいつも頭痛の種
個人的に２ヶ月に一回くらい踏む罠
jax.jit の中で Pythonの if を使って、しかも条件に 引数（bool）を入れると、JAXはそれを「トレース対象（Tracer）」として扱うことがあり、if が評価できず落ちる。
直し方（王道2択）
A) 条件分岐を jitの外に出す（おすすめ、確実）
B) jax.jit(... static_argnames=...) で 条件分岐を静的にする（やや玄人向け）
Flaxでトランスフォーマー組むときも、これで本当に泣きます。