import tensorflow as tf


def variational_expectations( Fmu, Fvar, phi, num_gauss_hermite_points=20):
    """
    Compute the expected value of a function phi, given a Gaussian
    distribution for the input values.
    if
        q(f) = N(Fmu, Fvar)
    then this method computes
       \int phi(f) q(f) df.
    Here, we implement a default Gauss-Hermite quadrature routine
    """
    gh_x, gh_w = hermgauss(num_gauss_hermite_points)
    gh_x = gh_x.reshape(1, -1)
    gh_w = gh_w.reshape(-1, 1) / np.sqrt(np.pi)
    shape = tf.shape(Fmu)
    Fmu, Fvar = [tf.reshape(e, (-1, 1)) for e in (Fmu, Fvar)]
    X = gh_x * tf.sqrt(2.0 * Fvar) + Fmu
    logp = phi(X)
    return tf.reshape(tf.matmul(logp, gh_w), shape)



import tensorflow as tf

def block_diagonal(matrices, dtype=tf.float32):
    """Constructs block-diagonal matrices from a list of batched 2D tensors.
    Args:
      matrices: A list of Tensors with shape [..., N_i, M_i] (i.e. a list of
        matrices with the same batch dimension).
      dtype: Data type to use. The Tensors in `matrices` must match this dtype.
