Returns a sample from a D-dimensional Multivariate Normal distribution
    :param mean: [..., N, D]
    :param cov: [..., N, D] or [..., N, D, 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 [..., (S), N, D], S = num_samples

def _sample_mvn(mean, cov, cov_structure=None, num_samples=None):
    """
    Returns a sample from a D-dimensional Multivariate Normal distribution
    :param mean: [..., N, D]
    :param cov: [..., N, D] or [..., N, D, 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 [..., (S), N, D], S = num_samples
    """
    mean_shape = tf.shape(mean)
    S = num_samples if num_samples is not None else 1
    D = mean_shape[-1]
    leading_dims = mean_shape[:-2]
    num_leading_dims = tf.size(leading_dims)

    if cov_structure == "diag":
        # mean: [..., N, D] and cov [..., N, D]
        with tf.control_dependencies([tf.assert_equal(tf.rank(mean), tf.rank(cov))]):
            eps_shape = tf.concat([leading_dims, [S], mean_shape[-2:]], 0)
            eps = tf.random_normal(eps_shape, dtype=settings.float_type)  # [..., S, N, D]
            samples = mean[..., None, :, :] + tf.sqrt(cov)[..., None, :, :] * eps  # [..., S, N, D]
    elif cov_structure == "full":
        # mean: [..., N, D] and cov [..., N, D, D]
        with tf.control_dependencies([tf.assert_equal(tf.rank(mean) + 1, tf.rank(cov))]):
            jittermat = (
                tf.eye(D, batch_shape=mean_shape[:-1], dtype=settings.float_type)
                * settings.jitter
            )  # [..., N, D, D]
            eps_shape = tf.concat([mean_shape, [S]], 0)
            eps = tf.random_normal(eps_shape, dtype=settings.float_type)  # [..., N, D, S]
            chol = tf.cholesky(cov + jittermat)  # [..., N, D, D]
            samples = mean[..., None] + tf.matmul(chol, eps)  # [..., N, D, S]
            samples = misc.leading_transpose(samples, [..., -1, -3, -2])  # [..., S, N, D]
    else:
        raise NotImplementedError  # pragma: no cover

    if num_samples is None:
        return samples[..., 0, :, :]  # [..., N, D]
    return samples