Repository: google/jax

Issue:
random.multivariate_normal broadcasting error  
There appears to be a broadcasting bug in jax.random.multivariate_normal for the SVD and eigh methods that
does not exist for the cholseky method. Example code with the resulting error message (for the SVD) case are appended below. 

```python
import jax.numpy as jnp
from jax import random

key=random.PRNGKey(42)
key, subkey =random.split(key)  # Jax random number generator key
process_dim = 100
nsamp = 10
mean = jnp.zeros((nsamp, process_dim))
covar = (jnp.ones(nsamp)[..., None, None])*(jnp.diag(jnp.ones(process_dim))[None, ...])

print('mean.shape = ', mean.shape)
print('covar.shape = ', covar.shape)
# Cholesky works as expected
norm_samp_cholesky = random.multivariate_normal(subkey, mean, covar, shape=(nsamp,), method='cholesky')
# SVD and eigh fail with broadcasting error message
norm_samp_svd = random.multivariate_normal(subkey, mean, covar, shape=(nsamp,), method='svd')
norm_samp_eigh = random.multivariate_normal(subkey, mean, covar, shape=(nsamp,), method='eigh')

```

/opt/miniconda3/envs/qso_fitting/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py in fn(x1, x2)
    683 def _maybe_bool_binop(numpy_fn, lax_fn, bool_lax_fn, lax_doc=False):
    684   def fn(x1, x2):
--> 685     x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
    686     return lax_fn(x1, x2) if x1.dtype != bool_ else bool_lax_fn(x1, x2)
    687   fn = jit(fn, inline=True)

/opt/miniconda3/envs/qso_fitting/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py in _promote_args(fun_name, *args)
    575   _check_arraylike(fun_name, *args)
    576   _check_no_float0s(fun_name, *args)
--> 577   return _promote_shapes(fun_name, *_promote_dtypes(*args))
    578 
    579 def _promote_args_inexact(fun_name, *args):

/opt/miniconda3/envs/qso_fitting/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py in _promote_shapes(fun_name, *args)
    494       if config.jax_numpy_rank_promotion != "allow":
    495         _rank_promotion_warning_or_error(fun_name, shapes)
--> 496       result_rank = len(lax.broadcast_shapes(*shapes))
    497       return [broadcast_to(arg, (1,) * (result_rank - len(shp)) + shp)
    498               for arg, shp in zip(args, shapes)]

    [... skipping hidden 2 frame]

/opt/miniconda3/envs/qso_fitting/lib/python3.9/site-packages/jax/_src/lax/lax.py in broadcast_shapes(*shapes)
     92   result_shape = _try_broadcast_shapes(shapes)
     93   if result_shape is None:
---> 94     raise ValueError("Incompatible shapes for broadcasting: {}"
     95                      .format(tuple(map(tuple, shapes))))
     96   return result_shape

ValueError: Incompatible shapes for broadcasting: ((10, 100, 100), (1, 10, 100))



Discussion:

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -550,10 +550,10 @@ def _multivariate_normal(key, mean, cov, shape, dtype, method) -> jnp.ndarray:
 
   if method == 'svd':
     (u, s, _) = svd(cov)
-    factor = u * jnp.sqrt(s)
+    factor = u * jnp.sqrt(s[..., None, :])
   elif method == 'eigh':
     (w, v) = eigh(cov)
-    factor = v * jnp.sqrt(w)
+    factor = v * jnp.sqrt(w[..., None, :])
   else: # 'cholesky'
     factor = cholesky(cov)
   normal_samples = normal(key, shape + mean.shape[-1:], dtype)


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -771,18 +771,19 @@ def testMultivariateNormal(self, dim, dtype, method):
       self._CheckKolmogorovSmirnovCDF(whitened.ravel(), scipy.stats.norm().cdf)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_dim={}_mean_batch_size={}_cov_batch_size={}_shape={}"\
-       .format(dim, mean_batch_size, cov_batch_size, shape),
+      {"testcase_name": "_dim={}_mean_batch_size={}_cov_batch_size={}_shape={}_method={}"\
+       .format(dim, mean_batch_size, cov_batch_size, shape, method),
        "dim": dim,
        "mean_batch_size": mean_batch_size,
        "cov_batch_size": cov_batch_size,
-       "shape": shape}
+       "shape": shape, "method": method}
       for dim in [1, 2, 4]
       for mean_batch_size in [(), (3,), (2, 3)]
       for cov_batch_size in [(), (3,), (2, 3)]
-      for shape in [(), (1,), (5,)]))
+      for shape in [(), (1,), (5,)]
+      for method in ['cholesky', 'svd', 'eigh']))
   def testMultivariateNormalShapes(self, dim, mean_batch_size, cov_batch_size,
-                                   shape):
+                                   shape, method):
     r = np.random.RandomState(0)
     key = self.seed_prng(0)
     eff_batch_size = mean_batch_size \
@@ -793,7 +794,7 @@ def testMultivariateNormalShapes(self, dim, mean_batch_size, cov_batch_size,
     cov += 1e-3 * np.eye(dim)
     shape = shape + eff_batch_size
     with jax.numpy_rank_promotion('allow'):
-      samples = random.multivariate_normal(key, mean, cov, shape=shape)
+      samples = random.multivariate_normal(key, mean, cov, shape=shape, method=method)
     assert samples.shape == shape + (dim,)
 
   def testMultivariateNormalCovariance(self):