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

brainpy / BrainPy

Brain Dynamics Programming in Python

https://brainpy.readthedocs.io/

GNU General Public License v3.0

491 stars 90 forks source link

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

Closed Routhleck closed 1 month ago

Routhleck commented 1 month ago

Because of the deprecation of from jax.experimental.host_callback import id_tap, replace jax.experimental.host_callback with jax.pure_callback. Also support jax==0.4.28:

The kind argument to jax.numpy.sort and jax.numpy.argsort is now removed. Use stable=True or stable=False instead.