Fix wrong call signatures for jnp.shape, jnp.size, and jnp.ndim

Passing lists or sequences to these functions is deprecated starting in jax-ml/jax#26641

PiperOrigin-RevId: 730489201

Author: Jake VanderPlas

distrax/_src/distributions/distribution_test.py

@@ -36,7 +36,7 @@ def log_prob(self, value):
   @property
   def event_shape(self):
     """Shape of the events."""
-    return jnp.shape([])
+    return np.shape([])
 
 
 class DummyMultivariateDist(distribution.Distribution):