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