Repository: google/jax

Issue:
Multivariate Normal
Hi all,

I noticed that sampling from predefined distributions works a bit different in jax than it does in numpy. 

Is it possible to use the jax.random module to sample from a multivariate normal distribution (with dense covariance matrix)? If not, is there any other way?

So far I only managed to sample from one with a unit variance, using jax.random.normal.

Cheers,

 


Discussion:
Thanks for the question!

The best way to sample from a multivariate Gaussian is to sample a vector of iid unit-variance zero-mean Gaussians (which you can do with `jax.random.normal`), and then perform an affine transformation to produce the mean and covariance structure you want.

For example, let's take the zero-mean case (since it's easy to add a mean vector in later). If x ~ N(0, I) is a vector of unit-variance zero-mean Gaussians, then y = A x has covariance E[yy']  = E[(A x)(A x)'] = E[A x x' A'] = A E[xx'] A' = AA'.

So if we had a target covariance Q and mean mu, the procedure would look something like
1. sample a standard Gaussian like v ~ N(0, I)
2. return a sample w = A v + mu, where A is a square root of Q such that AA' = Q.

See also [this section on Wikipedia](https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Drawing_values_from_the_distribution).

There are many matrix square roots (all related by an orthogonal transformation on the right), and any one will work. The standard one to use here would be the Cholesky square root, since (IIRC) it's the cheapest to compute in terms of FLOPs. That's usually written as L = chol(Q) where LL' = Q. (Another option is the symmetric square root, which basically looks like diagonalizing Q, taking the square root of the eigenvalues, and then multiplying the result back out. Fun fact, if you take the QR of the symmetric square root, then R' is the Cholesky!)

Since we have `vmap` we don't need to worry about handling batch dimensions by hand. Maybe a good API for a multivariate normal would be

```python
# Q is an (n, n) matrix
# mu is an (n,) vector
w = random.multivariate_normal(key, Q, mu)
```
Thanks for the answer! 

Sorry for the confusion. I was mostly interested to know whether there is a multivariate normal function already implemented in jax, since I couldn't find one.

Are there any plans to add more distributions (including a multivariate normal) to the jax.random? 

P.S. I ended up doing the Cholesky decomposition. 
In case you find it useful, we have a multivariate normal distribution in [numpyro](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/distributions/continuous.py#L559) that's basically just doing what Matt proposed above (the remaining code is just to comply with our distributions API). cc. @fehiepsi. 
In JAX, there is a pending PR #269 which does the job.
Thanks a lot! I think this answers my question.

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -240,12 +240,19 @@ def _random_bits(key, bit_width, shape):
 ### random samplers
 
 
-def _check_shape(name, shape):
+def _check_shape(name, shape, *param_shapes):
   try:
     shape = tuple(map(int, shape))
   except TypeError:
     msg = "{} requires a concrete tuple of integers as shape argument, got {}."
     raise ValueError(msg.format(name, shape))
+  if param_shapes:
+    shape_ = lax.broadcast_shapes(shape, *param_shapes)
+    if shape != shape_:
+      msg = ("{} parameter shapes must be broadcast-compatible with shape "
+             "argument, and the result of broadcasting the shapes must equal "
+             "the shape argument, but got result {} for shape argument {}.")
+      raise ValueError(msg.format(name, shape_, shape))
 
 
 def uniform(key, shape=(), dtype=onp.float64, minval=0., maxval=1.):
@@ -253,7 +260,8 @@ def uniform(key, shape=(), dtype=onp.float64, minval=0., maxval=1.):
 
   Args:
     key: a PRNGKey used as the random key.
-    shape: a tuple of nonnegative integers representing the shape.
+    shape: optional, a tuple of nonnegative integers representing the result
+      shape. Default ().
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
     minval: optional, a minimum (inclusive) value for the range (default 0).
@@ -302,7 +310,7 @@ def randint(key, shape, minval, maxval, dtype=onp.int64):
     shape: a tuple of nonnegative integers representing the shape.
     minval: int or array of ints broadcast-compatible with ``shape``, a minimum
       (inclusive) value for the range.
-    maxval: int or array of ints broadcast-compatible with  ``shape``, a maximum
+    maxval: int or array of ints broadcast-compatible with ``shape``, a maximum
       (exclusive) value for the range.
     dtype: optional, an int dtype for the returned values (default int64 if
       jax_enable_x64 is true, otherwise int32).
@@ -315,7 +323,7 @@ def randint(key, shape, minval, maxval, dtype=onp.int64):
 
 @partial(jit, static_argnums=(1, 4))
 def _randint(key, shape, minval, maxval, dtype):
-  _check_shape("randint", shape)
+  _check_shape("randint", shape, minval.shape, maxval.shape)
   if not onp.issubdtype(dtype, onp.integer):
     raise TypeError("randint only accepts integer dtypes.")
 
@@ -400,7 +408,8 @@ def normal(key, shape=(), dtype=onp.float64):
 
   Args:
     key: a PRNGKey used as the random key.
-    shape: a tuple of nonnegative integers representing the shape.
+    shape: optional, a tuple of nonnegative integers representing the result
+      shape. Default ().
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
@@ -419,63 +428,84 @@ def _normal(key, shape, dtype):
   return onp.array(onp.sqrt(2), dtype) * lax.erf_inv(u)
 
 
-def multivariate_normal(key, mean=0.0, cov=1.0, shape=(), dtype=onp.float64):
-  dtype = xla_bridge.canonicalize_dtype(dtype)
-  return _multivariate_normal(key, mean, cov, shape, dtype)
-
-@partial(jit, static_argnums=(3, 4))
-def _multivariate_normal(key, mean, cov, shape, dtype):
-  """Sample multivariate normal random values with given shape, mean, and covariance.
+def multivariate_normal(key, mean, cov, shape=None, dtype=onp.float64):
+  """Sample multivariate normal random values with given mean and covariance.
 
   Args:
     key: a PRNGKey used as the random key.
-    mean: optional, a scalar or array of mean values along each dimension
-    cov: optional, a scalar (isotropic), vector (diagonal covariance matrix), or full covariance matrix
-    shape: optional, a tuple of nonnegative integers representing the shape.
+    mean: a mean vector of shape ``(..., n)``.
+    cov: a positive definite covariance matrix of shape ``(..., n, n)``. The
+      batch shape ``...`` must be broadcast-compatible with that of ``mean``.
+    shape: optional, a tuple of nonnegative integers specifying the result
+      batch shape; that is, the prefix of the result shape excluding the last
+      axis. Must be broadcast-compatible with ``mean.shape[:-1]`` and
+      ``cov.shape[:-2]``. The default (None) produces a result batch shape by
+      broadcasting together the batch shapes of ``mean`` and ``cov``.
+    dtype: optional, a float dtype for the returned values (default float64 if
+      jax_enable_x64 is true, otherwise float32).
 
   Returns:
-    A random array with latent dimension of (max(asarray(mean).ndim, asarray(cov).ndim)),)
+    A random array with the specified dtype and shape given by
+    ``shape + mean.shape[-1:]`` if ``shape`` is not None, or else
+    ``broadcast_shapes(mean.shape[:-1], cov.shape[:-2]) + mean.shape[-1:]``.
   """
-  _check_shape("multivariate_normal", shape)
-  if hasattr(mean, "shape") and mean.ndim > 1:
-      raise ValueError("Mean cannot have more than 1 dimension.")
-  if hasattr(cov, "shape") and cov.ndim > 0:
-      if cov.ndim > 2:
-        raise ValueError("Covariance matrix cannot have more than 2 dimensions.")
-      shape = shape + cov.shape[:1]
-      normal_samples = normal(key, shape, dtype)
-      if cov.ndim == 2:
-        samples = np.tensordot(normal_samples, cholesky(cov), axes=1)
-      else:
-        samples = normal_samples * np.sqrt(cov)
+  dtype = xla_bridge.canonicalize_dtype(dtype)
+  return _multivariate_normal(key, mean, cov, shape, dtype)
+
+@partial(jit, static_argnums=(3, 4))
+def _multivariate_normal(key, mean, cov, shape, dtype):
+  if not onp.ndim(mean) >= 1:
+    msg = "multivariate_normal requires mean.ndim >= 1, got mean.ndim == {}"
+    raise ValueError(msg.format(onp.ndim(mean)))
+  if not onp.ndim(cov) >= 2:
+    msg = "multivariate_normal requires cov.ndim >= 2, got cov.ndim == {}"
+    raise ValueError(msg.format(onp.ndim(cov)))
+  n = mean.shape[-1]
+  if onp.shape(cov)[-2:] != (n, n):
+    msg = ("multivariate_normal requires cov.shape == (..., n, n) for n={n}, "
+           "but got cov.shape == {shape}.")
+    raise ValueError(msg.format(n=n, shape=onp.shape(cov)))
+
+  if shape is None:
+    shape = lax.broadcast_shapes(mean.shape[:-1], cov.shape[:-2])
   else:
-    if hasattr(mean, "shape") and mean.ndim > 0:
-      shape = shape + mean.shape[:1]
-    normal_samples = normal(key, shape, dtype)
-    samples = np.sqrt(cov) * normal_samples
-  return samples + mean
+    _check_shape("normal", shape, mean.shape[:-1], mean.shape[:-2])
 
+  chol_factor = cholesky(cov)
+  normal_samples = normal(key, shape + mean.shape[-1:], dtype)
+  return mean + np.tensordot(normal_samples, chol_factor, [-1, 1])
 
-def truncated_normal(key, lower, upper, shape=(), dtype=onp.float64):
+
+def truncated_normal(key, lower, upper, shape=None, dtype=onp.float64):
   """Sample truncated standard normal random values with given shape and dtype.
 
   Args:
     key: a PRNGKey used as the random key.
-    lower: a floating-point lower bound for truncation.
-    upper: a floating-point upper bound for truncation.
-    shape: a tuple of nonnegative integers representing the shape.
+    lower: a float or array of floats representing the lower bound for
+      truncation. Must be broadcast-compatible with ``upper``.
+    upper: a float or array of floats representing the  upper bound for
+      truncation. Must be broadcast-compatible with ``lower``.
+    shape: optional, a tuple of nonnegative integers specifying the result
+      shape. Must be broadcast-compatible with ``lower`` and ``upper``. The
+      default (None) produces a result shape by broadcasting ``lower`` and
+      ``upper``.
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
   Returns:
-    A random array with the specified shape and dtype.
+    A random array with the specified dtype and shape given by ``shape`` if
+    ``shape`` is not None, or else by broadcasting ``lower`` and ``upper``.
   """
   dtype = xla_bridge.canonicalize_dtype(dtype)
   return _truncated_normal(key, lower, upper, shape, dtype)
 
 @partial(jit, static_argnums=(3, 4))
 def _truncated_normal(key, lower, upper, shape, dtype):
-  _check_shape("truncated_normal", shape)
+  if shape is None:
+    shape = lax.broadcast_shapes(lower.shape, upper.shape)
+  else:
+    _check_shape("truncated_normal", shape, lower.shape, upper.shape)
+
   sqrt2 = onp.array(onp.sqrt(2), dtype)
   a = lax.erf(lax.convert_element_type(lower, dtype) / sqrt2)
   b = lax.erf(lax.convert_element_type(upper, dtype) / sqrt2)
@@ -485,18 +515,20 @@ def _truncated_normal(key, lower, upper, shape, dtype):
   return sqrt2 * lax.erf_inv(a + u * (b - a))
 
 
-def bernoulli(key, p=onp.float32(0.5), shape=()):
+def bernoulli(key, p=onp.float32(0.5), shape=None):
   """Sample Bernoulli random values with given shape and mean.
 
   Args:
     key: a PRNGKey used as the random key.
-    p: optional, an array-like of floating dtype broadcastable to `shape` for
-      the mean of the random variables (default 0.5).
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    p: optional, a float or array of floats for the mean of the random
+      variables. Must be broadcast-compatible with ``shape``. Default 0.5.
+    shape: optional, a tuple of nonnegative integers representing the result
+      shape. Must be broadcast-compatible with ``p.shape``. The default (None)
+      produces a result shape equal to ``p.shape``.
 
   Returns:
-    A random array with the specified shape and boolean dtype.
+    A random array with boolean dtype and shape given by ``shape`` if ``shape``
+    is not None, or else ``p.shape``.
   """
   dtype = xla_bridge.canonicalize_dtype(lax.dtype(p))
   if not onp.issubdtype(dtype, onp.floating):
@@ -507,39 +539,45 @@ def bernoulli(key, p=onp.float32(0.5), shape=()):
 
 @partial(jit, static_argnums=(2,))
 def _bernoulli(key, p, shape):
-  _check_shape("bernoulli", shape)
-  shape = shape or onp.shape(p)
-  if onp.shape(p) != shape:
-    p = np.broadcast_to(p, shape)
-  return lax.lt(uniform(key, shape, lax.dtype(p)), p)
+  if shape is None:
+    shape = p.shape
+  else:
+    _check_shape("bernoulli", shape, p.shape)
 
+  return uniform(key, shape, lax.dtype(p)) < p
 
-def beta(key, a, b, shape=(), dtype=onp.float64):
+
+def beta(key, a, b, shape=None, dtype=onp.float64):
   """Sample Bernoulli random values with given shape and mean.
 
   Args:
     key: a PRNGKey used as the random key.
-    a: an array-like broadcastable to `shape` and used as the shape parameter
-      alpha of the random variables.
-    b: an array-like broadcastable to `shape` and used as the shape parameter
-      beta of the random variables.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    a: a float or array of floats broadcast-compatible with ``shape``
+      representing the first parameter "alpha".
+    b: a float or array of floats broadcast-compatible with ``shape``
+      representing the second parameter "beta".
+    shape: optional, a tuple of nonnegative integers specifying the result
+      shape. Must be broadcast-compatible with ``a`` and ``b``. The default
+      (None) produces a result shape by broadcasting ``a`` and ``b``.
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
   Returns:
-    A random array with the specified shape and dtype.
+    A random array with the specified dtype and shape given by ``shape`` if
+    ``shape`` is not None, or else by broadcasting ``a`` and ``b``.
   """
   dtype = xla_bridge.canonicalize_dtype(dtype)
   return _beta(key, a, b, shape, dtype)
 
 @partial(jit, static_argnums=(3, 4))
 def _beta(key, a, b, shape, dtype):
-  _check_shape("beta", shape)
+  if shape is None:
+    shape = lax.broadcast_shapes(a.shape, b.shape)
+  else:
+    _check_shape("beta", shape, a.shape, b.shape)
+
   a = lax.convert_element_type(a, dtype)
   b = lax.convert_element_type(b, dtype)
-  shape = shape or lax.broadcast_shapes(np.shape(a), np.shape(b))
   key_a, key_b = split(key)
   gamma_a = gamma(key_a, a, shape, dtype)
   gamma_b = gamma(key_b, b, shape, dtype)
@@ -551,8 +589,8 @@ def cauchy(key, shape=(), dtype=onp.float64):
 
   Args:
     key: a PRNGKey used as the random key.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    shape: optional, a tuple of nonnegative integers representing the result
+      shape. Default ().
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
@@ -570,29 +608,41 @@ def _cauchy(key, shape, dtype):
   return lax.tan(lax.mul(pi, lax.sub(u, _constant_like(u, 0.5))))
 
 
-def dirichlet(key, alpha, shape=(), dtype=onp.float64):
+def dirichlet(key, alpha, shape=None, dtype=onp.float64):
   """Sample Cauchy random values with given shape and float dtype.
 
   Args:
     key: a PRNGKey used as the random key.
-    alpha: an array-like with `alpha.shape[:-1]` broadcastable to `shape` and
-      used as the concentration parameter of the random variables.
-    shape: optional, a tuple of nonnegative integers representing the batch
-      shape (defaults to `alpha.shape[:-1]`).
+    alpha: an array of shape ``(..., n)`` used as the concentration
+      parameter of the random variables.
+    shape: optional, a tuple of nonnegative integers specifying the result
+      batch shape; that is, the prefix of the result shape excluding the last
+      element of value ``n``. Must be broadcast-compatible with
+      ``alpha.shape[:-1]``. The default (None) produces a result shape equal to
+      ``alpha.shape``.
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
   Returns:
-    A random array with the specified shape and dtype.
+    A random array with the specified dtype and shape given by
+    ``shape + (alpha.shape[-1],)`` if ``shape`` is not None, or else
+    ``alpha.shape``.
   """
   dtype = xla_bridge.canonicalize_dtype(dtype)
   return _dirichlet(key, alpha, shape, dtype)
 
 @partial(jit, static_argnums=(2, 3))
 def _dirichlet(key, alpha, shape, dtype):
-  _check_shape("dirichlet", shape)
-  alpha = asarray(alpha, dtype)
-  shape = shape or alpha.shape[:-1]
+  if not onp.ndim(alpha) >= 1:
+    msg = "dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}"
+    raise ValueError(msg.format(onp.ndim(alpha)))
+
+  if shape is None:
+    shape = alpha.shape[:-1]
+  else:
+    _check_shape("dirichlet", shape, alpha.shape[:-1])
+
+  alpha = lax.convert_element_type(alpha, dtype)
   gamma_samples = gamma(key, alpha, shape + alpha.shape[-1:], dtype)
   return gamma_samples / np.sum(gamma_samples, axis=-1, keepdims=True)
 
@@ -602,8 +652,8 @@ def exponential(key, shape=(), dtype=onp.float64):
 
   Args:
     key: a PRNGKey used as the random key.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    shape: optional, a tuple of nonnegative integers representing the result
+      shape. Default ().
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
@@ -676,7 +726,6 @@ def _next_kxv(kxv):
   z = lax.mul(lax.mul(d, V), boost)
   return lax.select(lax.eq(z, zero), onp.finfo(z.dtype).tiny, z)
 
-
 _bivariate_coef = [[0.16009398, -0.094634816, 0.025146379, -0.0030648348,
                     1, 0.3266811, 0.10406087, 0.0014179033],
                    [0.53487893, 0.12980707, 0.06573594, -0.0015649787,
@@ -684,7 +733,6 @@ def _next_kxv(kxv):
                    [0.040121005, -0.0065914079, -0.002628604, -0.0013441777,
                     0.017050642, -0.0021309345, 0.00085092385, -1.5248239e-07]]
 
-
 def _gamma_grad_one(z, alpha):
     # Ref 1: Pathwise Derivatives Beyond the Reparameterization Trick, Martin & Fritz
     # Ref 2: Case 4 follows https://github.com/fritzo/notebooks/blob/master/gamma-reparameterized.ipynb
@@ -780,14 +828,12 @@ def _case4(zagf):
     _, _, grad, flag = lax.while_loop(lambda zagf: ~zagf[3], _case4, (z, alpha, grad, flag))
     return grad
 
-
 def _gamma_grad(sample, a):
     samples = np.reshape(sample, -1)
     alphas = np.reshape(a, -1)
     grads = vmap(_gamma_grad_one)(samples, alphas)
     return grads.reshape(a.shape)
 
-
 @custom_transforms
 def _gamma_impl(key, a):
     alphas = np.reshape(a, -1)
@@ -795,34 +841,37 @@ def _gamma_impl(key, a):
     samples = vmap(_gamma_one)(keys, alphas)
     return np.reshape(samples, np.shape(a))
 
-
 defjvp(_gamma_impl, None,
        lambda tangent, ans, key, a, **kwargs: tangent * _gamma_grad(ans, a))
 
-
-def gamma(key, a, shape=(), dtype=onp.float64):
+def gamma(key, a, shape=None, dtype=onp.float64):
   """Sample Gamma random values with given shape and float dtype.
 
   Args:
     key: a PRNGKey used as the random key.
-    a: an array-like broadcastable to `shape` and used as the shape parameter
-      of the random variables.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    a: a float or array of floats broadcast-compatible with ``shape``
+      representing the parameter of the distribution.
+    shape: optional, a tuple of nonnegative integers specifying the result
+      shape. Must be broadcast-compatible with ``a``. The default (None)
+      produces a result shape equal to ``a.shape``.
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
   Returns:
-    A random array with the specified shape and dtype.
+    A random array with the specified dtype and with shape given by ``shape`` if
+    ``shape`` is not None, or else by ``a.shape``.
   """
   dtype = xla_bridge.canonicalize_dtype(dtype)
   return _gamma(key, a, shape, dtype)
 
 @partial(jit, static_argnums=(2, 3))
 def _gamma(key, a, shape, dtype):
-  _check_shape("gamma", shape)
+  if shape is None:
+    shape = a.shape
+  else:
+    _check_shape("gamma", shape, a.shape)
+
   a = lax.convert_element_type(a, dtype)
-  shape = shape or onp.shape(a)
   if onp.shape(a) != shape:
     a = np.broadcast_to(a, shape)
   return _gamma_impl(key, a)
@@ -833,8 +882,8 @@ def gumbel(key, shape=(), dtype=onp.float64):
 
   Args:
     key: a PRNGKey used as the random key.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    shape: optional, a tuple of nonnegative integers representing the result
+      shape. Default ().
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
@@ -856,8 +905,8 @@ def laplace(key, shape=(), dtype=onp.float64):
 
   Args:
     key: a PRNGKey used as the random key.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    shape: optional, a tuple of nonnegative integers representing the result
+      shape. Default ().
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
@@ -880,8 +929,8 @@ def logistic(key, shape=(), dtype=onp.float64):
 
   Args:
     key: a PRNGKey used as the random key.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    shape: optional, a tuple of nonnegative integers representing the result
+      shape. Default ().
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
@@ -897,33 +946,36 @@ def _logistic(key, shape, dtype):
   return logit(uniform(key, shape, dtype))
 
 
-def pareto(key, b, shape=(), dtype=onp.float64):
+def pareto(key, b, shape=None, dtype=onp.float64):
   """Sample Pareto random values with given shape and float dtype.
 
   Args:
     key: a PRNGKey used as the random key.
-    b: an array-like broadcastable to `shape` and used as the shape parameter
-      of the random variables.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    a: a float or array of floats broadcast-compatible with ``shape``
+      representing the parameter of the distribution.
+    shape: optional, a tuple of nonnegative integers specifying the result
+      shape. Must be broadcast-compatible with ``b``. The default (None)
+      produces a result shape equal to ``b.shape``.
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
   Returns:
-    A random array with the specified shape and dtype.
+    A random array with the specified dtype and with shape given by ``shape`` if
+    ``shape`` is not None, or else by ``b.shape``.
   """
   dtype = xla_bridge.canonicalize_dtype(dtype)
   return _pareto(key, b, shape, dtype)
 
 @partial(jit, static_argnums=(2, 3))
 def _pareto(key, b, shape, dtype):
-  _check_shape("pareto", shape)
+  if shape is None:
+    shape = b.shape
+  else:
+    _check_shape("pareto", shape)
+
   b = lax.convert_element_type(b, dtype)
-  shape = shape or onp.shape(b)
-  if onp.shape(b) != shape:
-    b = np.broadcast_to(b, shape)
   e = exponential(key, shape, dtype)
-  return lax.exp(lax.div(e, b))
+  return lax.exp(e / b)
 
 
 def t(key, df, shape=(), dtype=onp.float64):
@@ -931,24 +983,29 @@ def t(key, df, shape=(), dtype=onp.float64):
 
   Args:
     key: a PRNGKey used as the random key.
-    df: an array-like broadcastable to `shape` and used as the shape parameter
-      of the random variables.
-    shape: optional, a tuple of nonnegative integers representing the shape
-      (default scalar).
+    df: a float or array of floats broadcast-compatible with ``shape``
+      representing the parameter of the distribution.
+    shape: optional, a tuple of nonnegative integers specifying the result
+      shape. Must be broadcast-compatible with ``df``. The default (None)
+      produces a result shape equal to ``df.shape``.
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
 
   Returns:
-    A random array with the specified shape and dtype.
+    A random array with the specified dtype and with shape given by ``shape`` if
+    ``shape`` is not None, or else by ``df.shape``.
   """
   dtype = xla_bridge.canonicalize_dtype(dtype)
   return _t(key, df, shape, dtype)
 
 @partial(jit, static_argnums=(2, 3))
 def _t(key, df, shape, dtype):
-  _check_shape("t", shape)
+  if shape is None:
+    shape = df.shape
+  else:
+    _check_shape("t", shape, df.shape)
+
   df = lax.convert_element_type(df, dtype)
-  shape = shape or onp.shape(df)
   key_n, key_g = split(key)
   n = normal(key_n, shape, dtype)
   two = _constant_like(n, 2)


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -16,8 +16,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from functools import partial
 from unittest import SkipTest
-import re
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -218,7 +218,9 @@ def testCauchy(self, dtype):
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_alpha={}_{}".format(alpha, dtype),
        "alpha": alpha, "dtype": onp.dtype(dtype).name}
-      for alpha in [[0.2, 1., 5.]]
+      for alpha in [
+          onp.array([0.2, 1., 5.]),
+      ]
       for dtype in [onp.float32, onp.float64]))
   def testDirichlet(self, alpha, dtype):
     key = random.PRNGKey(0)
@@ -276,7 +278,7 @@ def testGammaGrad(self, alpha):
     rng = random.PRNGKey(0)
     alphas = onp.full((100,), alpha)
     z = random.gamma(rng, alphas)
-    actual_grad = api.grad(lambda x: (random.gamma(rng, x)).sum())(alphas)
+    actual_grad = api.grad(lambda x: random.gamma(rng, x).sum())(alphas)
 
     eps = 0.01 * alpha / (1.0 + onp.sqrt(alpha))
     cdf_dot = (scipy.stats.gamma.cdf(z, alpha + eps)
@@ -367,35 +369,33 @@ def testT(self, df, dtype):
       self._CheckKolmogorovSmirnovCDF(samples, scipy.stats.t(df).cdf)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_mean={}_cov={}_{}".format(mean, cov, dtype),
-        "mean": mean, "cov": cov, "dtype": dtype}
-      for mean in [0, 5, onp.asarray([1, -2, 3]), onp.asarray([[1]])]
-      for cov in [.1, 5, onp.asarray([4, 5, 6]),
-        onp.asarray([[4.60, 2.86, 2.33],
-        [2.86, 3.04, 1.74],
-        [2.33, 1.74, 1.83]]),
-        onp.asarray([[[1]]])]
+      {"testcase_name": "_{}D_{}".format(dim, onp.dtype(dtype).name),
+       "dim": dim, "dtype": dtype}
+      for dim in [1, 3, 5]
       for dtype in [onp.float32, onp.float64]))
-  def testMultivariateNormal(self, mean, cov, dtype):
+  def testMultivariateNormal(self, dim, dtype):
+    r = onp.random.RandomState(dim)
+    mean = r.randn(dim)
+    cov_factor = r.randn(dim, dim)
+    cov = onp.dot(cov_factor, cov_factor.T) + dim * onp.eye(dim)
+
     key = random.PRNGKey(0)
-    rand = lambda key, mean, cov: random.multivariate_normal(key, mean, cov, (1000,), dtype)
+    rand = partial(random.multivariate_normal, mean=mean, cov=cov,
+                   shape=(10000,))
     crand = api.jit(rand)
-    if hasattr(cov, "shape") and cov.ndim > 2 or hasattr(mean, "shape") and mean.ndim > 1:
-      self.assertRaises(ValueError, lambda: rand(key, mean, cov))
-      self.assertRaises(ValueError, lambda: crand(key, mean, cov))
-      return
-
-    uncompiled_samples = rand(key, mean, cov)
-    compiled_samples = crand(key, mean, cov)
-    if hasattr(cov, "shape") and cov.ndim == 2:
-      inv_scale = scipy.linalg.lapack.dtrtri(onp.linalg.cholesky(cov), lower=True)[0]
-      rescale = lambda x: onp.tensordot(x, inv_scale, axes=(-1, 1))
-    else:
-      rescale = lambda x: x / np.sqrt(cov)
+
+    uncompiled_samples = onp.asarray(rand(key), onp.float64)
+    compiled_samples = onp.asarray(crand(key), onp.float64)
+
+    inv_scale = scipy.linalg.lapack.dtrtri(onp.linalg.cholesky(cov), lower=True)[0]
     for samples in [uncompiled_samples, compiled_samples]:
-      self._CheckKolmogorovSmirnovCDF(
-          rescale(samples - mean).reshape(-1),
-          scipy.stats.norm().cdf)
+      centered = samples - mean
+      whitened = onp.einsum('nj,ij->ni', centered, inv_scale)
+
+      # This is a quick-and-dirty multivariate normality check that tests that a
+      # uniform mixture of the marginals along the covariance matrix's
+      # eigenvectors follow a standard normal distribution.
+      self._CheckKolmogorovSmirnovCDF(whitened.ravel(), scipy.stats.norm().cdf)
 
   def testIssue222(self):
     x = random.randint(random.PRNGKey(10003), (), 0, 0)
@@ -417,7 +417,7 @@ def feature_map(n, d, sigma=1.0, seed=123):
       phi = lambda x, t: np.sqrt(2.0 / d) * np.cos(np.matmul(W, x) + w*t + b)
       return phi
 
-    self.assertRaisesRegex(ValueError, re.compile(r'.*requires a concrete.*'),
+    self.assertRaisesRegex(ValueError, '.*requires a concrete.*',
                            lambda: feature_map(5, 3))
 
   def testIssue756(self):