Closed Routhleck closed 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
from jax.experimental.host_callback import id_tap
jax.experimental.host_callback
jax.pure_callback
Because of the deprecation of
from jax.experimental.host_callback import id_tap
, replacejax.experimental.host_callback
withjax.pure_callback