Remove references to deprecated jax.ShapedArray

This is deprecated as of https://github.com/google/jax/pull/15263: most users will never need to use ShapedArray directly, and so having it exposed in the top-level public namespace causes undue confusion.

PiperOrigin-RevId: 520846838

kfac_ferminet_alpha/tag_graph_matcher.py

@@ -154,7 +154,7 @@ def write_env(env, var, val):
 
 def abstract_single_value(value):
   if isinstance(value, jnp.ndarray):
-    value = jax.ShapedArray(np.shape(value), np.result_type(value))
+    value = jax_core.ShapedArray(np.shape(value), np.result_type(value))
     return pe.PartialVal.unknown(value)
   else:
     return value