import tensorflow as tf

    """
    Returns a sample from a D-dimensional Multivariate Normal distribution
    :param mean: N x D
    :param cov: N x D or N x D x D
    :param cov_structure: "diag" or "full"
    - "diag": cov holds the diagonal elements of the covariance matrix
    - "full": cov holds the full covariance matrix (without jitter)
    :return: sample from the MVN of shape N x D
    """
    eps = tf.random_normal(tf.shape(mean), dtype=settings.float_type)  # N x P
    if cov_structure == "diag":
        sample = mean + tf.sqrt(cov) * eps  # N x P
    elif cov_structure == "full":
        cov = cov + (tf.eye(tf.shape(mean)[1], dtype=settings.float_type) * settings.numerics.jitter_level)[None, ...]  # N x P x P
        chol = tf.cholesky(cov)  # N x P x P
        return mean + (tf.matmul(chol, eps[..., None])[..., 0])  # N x P
    else:
        raise NotImplementedError  # pragma: no cover
