Replace deprecated jax.core.get_aval with jax.typeof

`core.get_aval` is deprecated as of JAX v0.8.2 and will be removed in a future release.

PiperOrigin-RevId: 840774897

Author: Jake VanderPlas

distrax/_src/utils/math.py

@@ -19,7 +19,6 @@
 
 import chex
 import jax
-from jax import core as jax_core
 from jax.custom_derivatives import SymbolicZero
 import jax.numpy as jnp
 
@@ -72,7 +71,7 @@ def multiply_no_nan_jvp(
   x, y = primals
   x_dot, y_dot = tangents
   primal_out = multiply_no_nan(x, y)
-  primal_aval = jax_core.get_aval(primal_out)
+  primal_aval = jax.typeof(primal_out)
   result_aval = primal_aval.at_least_vspace()
   tangent_out_1 = scale_maybe_symbolic(result_aval, x_dot, y)
   tangent_out_2 = scale_maybe_symbolic(result_aval, y_dot, x)