Future-proof calls to jnp.solve on batched 1D inputs.

This has been deprecated since JAX v0.4.25, and is no longer supported in JAX v0.5.0. PiperOrigin-RevId: 712980403

Author: Jake VanderPlas

trax/layers/core.py

@@ -808,8 +808,7 @@ def log_gaussian_pdf(x, mu, sigma): # pylint: disable=invalid-name
  """
  a = mu.shape[-1] * jnp.log(2 * jnp.pi)
  _, b = jnp.linalg.slogdet(sigma)
- y = jnp.linalg.solve(sigma, x - mu)
- y = jnp.expand_dims(y, axis=-1)
+ y = jnp.linalg.solve(sigma, (x - mu)[..., None])
  xm = jnp.expand_dims(x - mu, axis=-2)
  c = jnp.matmul(xm, y)
  c = jnp.squeeze(jnp.squeeze(c, axis=-1), axis=-1)