def sample(self, rng: jax.Array, n_samples: int = 1) -> jnp.ndarray:
        """
        Sample from the multivariate Gaussian.

        :param rng: jax.Array
            Random number generator.
        :param n_samples: int
            Number of samples.

        :return: Array
            Samples. `shape = (number of samples, dimension)`
        """
        keys = random.split(rng, n_samples)
        return vmap(lambda key: random.multivariate_normal(key, self.mean, self.cov))(
            keys
        )

        Sample from the multivariate Gaussian.

        :param rng: jax.Array
            Random number generator.
        :param n_samples: int
            Number of samples.

        :return: Array
            Samples. `shape = (number of samples, dimension)`
        
fortuna/distribution/gaussian.py