I am getting the same error by calling jax.random.multivariate_normal . Please fix this, it is very basic functionality.