[fix] Replace jax.experimental.host_callback with jax.pure_callback#670

Merged

chaoming0625 merged 4 commits intomasterfrom pure_callback

May 14, 2024

+144-117

Commits on May 12, 2024

Revert "fix issue #661 (#662 )"
Routhleck
committed
Replace
Routhleck
committed

Commits on May 13, 2024

Support jax==0.4.28
Routhleck
committed

Commits on May 14, 2024

Fix JIT bugs and Replace deprecated functions
Routhleck
committed