Use jax.lib.xla_client instead of jax._src.lib.xla_client to use the public API.

PiperOrigin-RevId: 675686360

Author: lukebaumann

pathwaysutils/persistence/helper.py

@@ -20,7 +20,7 @@
 
 import jax
 from jax import core
-from jax._src.lib import xla_client as xc
+from jax.lib import xla_client as xc
 import numpy as np
 from pathwaysutils import plugin_executable