Migrate references to jax.lib.xla_extension to jax.extend.

PiperOrigin-RevId: 789636726

Author: danielsuo

pathwaysutils/proxy_backend.py

@@ -16,7 +16,6 @@
 import jax
 from jax.extend import backend
 from jax.lib.xla_extension import ifrt_proxy
-from jax.lib.xla_extension import ifrt_proxy
 
 
 def register_backend_factory():

pathwaysutils/test/proxy_backend_test.py

@@ -18,9 +18,9 @@
 import jax
 from jax.extend import backend
 from jax.lib.xla_extension import ifrt_proxy
-from jax.lib.xla_extension import ifrt_proxy
 from pathwaysutils import proxy_backend
 
+
 from absl.testing import absltest