Fixes and tests for jax.random.multivariate_normal (#4002)
* Fix bug #3997, change `jax.random.multivariate_normal` to handle batches of covariance matrices. It works as long as mean and covariance are broadcast-compatible, as specified in the docstring.
* Fix bug in multivariate_normal shape checking
Minor bug: should be checking for compatibility of `shape`, `mean`, and the the last two dimensions of the _covariance_ matrix.
* Add test for multivariate_normal shapes
This test checks that `jax.random.multivariate_normal` produces the expected output shape for various combinations of event dimension and `mean`, `covariance`, and `shape` shapes.
* Fix linter issues in tests/random_test.py
Trimming trialing whitespace and 80 char limit.
* Really trimming whitespace in tests/random_test.py
Arg. Have to fix my editor to do this automatically.