Replace jax.lib.xla_client._xla
with jax.lib.xla_extension
#1561
Loading
jax.lib.xla_client._xla
with jax.lib.xla_extension
#1561