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

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

        :return: Array
            Samples. `shape = (number of samples, dimension)`
        """
        _, key = random.split(rng)
        return self.mean + self.std * random.normal(key, (n_samples, self.dim))

        Sample from the diagonal 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