docs: 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