Refine JIT wrappers for new JAX for comaptiblity with jax>=0.8.2#809
Merged
chaoming0625 merged 3 commits intomasterfrom Dec 25, 2025
Merged
Refine JIT wrappers for new JAX for comaptiblity with `jax>=0.8.2`#809chaoming0625 merged 3 commits intomasterfrom
chaoming0625 merged 3 commits intomasterfrom