Replace usage of xla.abstractify with core.get_aval where feasible.

PiperOrigin-RevId: 730598559

Author: Jake VanderPlas

distrax/_src/utils/jittable.py

@@ -61,9 +61,9 @@ def _is_jax_data(x):
   if isinstance(x, (bool, int, float)) or x is None:
     return False
 
-  # Otherwise, try to make it into a tracer. If it succeeds, then it's JAX data.
+  # Otherwise, attempt to trace the value. If it succeeds, then it's JAX data.
   try:
-    jax.interpreters.xla.abstractify(x)
+    jax.eval_shape(lambda x: x, x)
     return True
   except TypeError:
     return False