File: gaussian_markov_chain.py

package info (click to toggle)
python-bayespy 0.6.2-1
links: PTS, VCS
area: main
in suites: forky, sid, trixie
size: 2,132 kB
sloc: python: 22,402; makefile: 156
file content (2098 lines) | stat: -rw-r--r-- 73,623 bytes
################################################################################
# Copyright (C) 2012-2014 Jaakko Luttinen
#
# This file is licensed under the MIT License.
################################################################################


"""
This module contains VMP nodes for Gaussian Markov chains.
"""

import numpy as np
import scipy

from bayespy.utils import misc
from bayespy.utils import linalg

from .node import Node, message_sum_multiply
from .deterministic import Deterministic
from .expfamily import ExponentialFamily
from .expfamily import ExponentialFamilyDistribution
from .expfamily import useconstructor
from .gaussian import (Gaussian,
                       GaussianMoments,
                       GaussianWishartMoments,
                       GaussianGammaMoments,
                       WrapToGaussianGamma,
                       WrapToGaussianWishart)
from .wishart import Wishart, WishartMoments
from .gamma import Gamma, GammaMoments
from .categorical import CategoricalMoments
from .node import Moments, ensureparents


class GaussianMarkovChainMoments(Moments):


    def __init__(self, N, D):
        self.N = N
        self.D = D
        return super().__init__()


    def compute_fixed_moments(self, x):
        u0 = x
        u1 = x[...,:,np.newaxis] * x[...,np.newaxis,:]
        u2 = x[...,:-1,:,np.newaxis] * x[...,1:,np.newaxis,:]
        return [u0, u1, u2]


    def rotate(self, u, R, logdet=None):

        if logdet is None:
            logdet = np.linalg.slogdet(R)[1]

        N = np.shape(u[0])[-2]

        # Transform moments and g
        u0 = linalg.mvdot(R, u[0])
        u1 = linalg.dot(R, u[1], R.T)
        u2 = linalg.dot(R, u[2], R.T)
        u = [u0, u1, u2]
        dg = -N * logdet

        return (u, dg)


class TemplateGaussianMarkovChainDistribution(ExponentialFamilyDistribution):
    """
    Sub-classes implement distribution specific computations.
    """


    def __init__(self, N, D):
        self.N = N
        self.D = D
        self.moments = GaussianMarkovChainMoments(N, D)
        super().__init__()

    def compute_message_to_parent(self, parent, index, u_self, *u_parents):
        raise NotImplementedError()

    def compute_weights_to_parent(self, index, weights):
        raise NotImplementedError()

    def compute_phi_from_parents(self, *u_parents, mask=True):
        raise NotImplementedError()

    def compute_moments_and_cgf(self, phi, mask=True):
        """
        Compute the moments and the cumulant-generating function.

        This basically performs the filtering and smoothing for the variable.

        Parameters
        ----------
        phi

        Returns
        -------
        u
        g
        """

        # Solve the Kalman filtering and smoothing problem
        y = phi[0]
        A = -2*phi[1]
        # Don't multiply phi[2] by two because it is a sum of the super- and
        # sub-diagonal blocks so we would need to divide by two anyway.
        B = -phi[2]

        (CovXnXn, CovXpXn, Xn, ldet) = linalg.block_banded_solve(A, B, y)

        # Compute moments
        u0 = Xn
        u1 = CovXnXn + Xn[...,:,np.newaxis] * Xn[...,np.newaxis,:]
        u2 = CovXpXn + Xn[...,:-1,:,np.newaxis] * Xn[...,1:,np.newaxis,:]
        u = [u0, u1, u2]

        # Compute cumulant-generating function
        g = -0.5 * np.einsum('...ij,...ij', u[0], phi[0]) + 0.5*ldet

        return (u, g)

    def compute_cgf_from_parents(self, *u_parents):
        raise NotImplementedError()

    def compute_fixed_moments_and_f(self, x, mask=True):
        """
        Compute u(x) and f(x) for given x.
        """
        u0 = x
        u1 = x[...,:,np.newaxis] * x[...,np.newaxis,:]
        u2 = x[...,:-1,:,np.newaxis] * x[...,1:,np.newaxis,:]
        u = [u0, u1, u2]

        f = -0.5 * np.shape(x)[-2] * np.shape(x)[-1] * np.log(2*np.pi)
        return (u, f)

    def plates_to_parent(self, index, plates):
        """
        Computes the plates of this node with respect to a parent.

        Child classes must implement this.

        Parameters
        -----------
        index : int
            The index of the parent node to use.
        """
        raise NotImplementedError()

    def plates_from_parent(self, index, plates):
        """
        Compute the plates using information of a parent node.

        Child classes must implement this.

        Parameters
        ----------
        index : int
            Index of the parent to use.
        """
        raise NotImplementedError()


    def rotate(self, u, phi, R, inv=None, logdet=None):

        (u, dg) = self.moments.rotate(u, R, logdet=logdet)

        # It would be more efficient and simpler, if you just rotated the
        # moments and didn't touch phi. However, then you would need to call
        # update() before lower_bound_contribution. This is more error-safe.

        if inv is None:
            inv = np.linalg.inv(R)

        # Transform parameters
        phi0 = linalg.mvdot(inv.T, phi[0])
        phi1 = linalg.dot(inv.T, phi[1], inv)
        phi2 = linalg.dot(inv.T, phi[2], inv)
        phi = [phi0, phi1, phi2]

        return (u, phi, dg)


    def compute_rotation_bound(self, u, u_mu_Lambda, u_A_V, R, inv=None, logdet=None):

        (Lambda_mu, Lambda_mumu, Lambda, logdetLambda) = u_mu_Lambda
        (V_A, V_AA, V, logdetV) = u_A_V

        V = misc.make_diag(V, ndim=1)

        R_XnXn = linalg.dot(R, self.XnXn)
        R_XpXp = linalg.dot(R, self.XpXp)
        R_X0X0 = linalg.dot(R, self.X0X0)

        tracedot(dot(Lambda, R_X0X0), R.T)
        tracedot(dot(V, R_XnXn), R.T)
        tracedot(dot(V_AA, R_XpXp), R.T)
        tracedot(dot(V_A, R_XpXn), R.T)
        (N - 1) * logdetV
        2 * N * logdetR

        logp = random.gaussian_logpdf(
            Lambda_R_X0X0_R + V_R_XnXn_R,
            V_A_R_XpXn_R,
            V_AA_R_XpXp_R,
            (N - 1) * logdetV + 2 * N * logdetR
        )

        logH = random.gaussian_entropy(
            -2 * M * logdetR,
            0
        )

        dlogp

        dlogH

        return (L, dL)


class _TemplateGaussianMarkovChain(ExponentialFamily):
    r"""
    VMP abstract node for Gaussian Markov chain.

    This is a general base class for different Gaussian Markov chain nodes.

    Output is Gaussian variables with mean, covariance and one-step
    cross-covariance.

    self.phi and self.u are defined in a particular way but otherwise the parent
    nodes may vary.

    Child classes must implement the following methods:
        _plates_to_parent
        _plates_from_parent

    See also
    --------
    bayespy.inference.vmp.nodes.gaussian.Gaussian
    bayespy.inference.vmp.nodes.wishart.Wishart

    """

    def random(self, *phi, plates=None):
        raise NotImplementedError()


def _compute_cgf_for_gaussian_markov_chain(mumu_Lambda, logdet_Lambda,
                                           logdet_nu, N):
    """
    Compute CGF using the moments of the parents.
    """

    g0 = -0.5 * mumu_Lambda #np.einsum('...ij,...ij->...', mumu, Lambda)

    g1 = 0.5 * logdet_Lambda
    if np.ndim(logdet_nu) == 1:
        g1 = g1 + 0.5 * (N-1) * np.sum(logdet_nu, axis=-1)
    elif np.shape(logdet_nu)[-2] == 1:
        g1 = g1 + 0.5 * (N-1) * np.sum(logdet_nu, axis=(-1,-2))
    else:
        g1 = g1 + 0.5 * np.sum(logdet_nu, axis=(-1,-2))

    return g0 + g1


class GaussianMarkovChainDistribution(TemplateGaussianMarkovChainDistribution):
    r"""
    Implementation of VMP formulas for Gaussian Markov chain

    The log probability density function of the prior:

    .. todo::  Fix inputs and their weight matrix in the equations.

    .. math::

       \log p(\mathbf{X} | \boldsymbol{\mu}, \mathbf{\Lambda},
       \mathbf{A}, \mathbf{B}, \boldsymbol{\nu})
       =& \log \mathcal{N}(\mathbf{x}_0|\boldsymbol{\mu}, \mathbf{\Lambda})
       + \sum^N_{n=1} \log \mathcal{N}(
         \mathbf{x}_n | \mathbf{Ax}_{n-1} + \mathbf{Bu}_n,
                        \mathrm{diag}(\boldsymbol{\nu}))
       \\
       =&
       - \frac{1}{2} \mathbf{x}_0^T \mathbf{\Lambda} \mathbf{x}_0
       + \frac{1}{2} \mathbf{x}_0^T \mathbf{\Lambda} \boldsymbol{\mu}
       + \frac{1}{2} \boldsymbol{\mu}^T \mathbf{\Lambda} \mathbf{x}_0
       - \frac{1}{2} \boldsymbol{\mu}^T \mathbf{\Lambda} \boldsymbol{\mu}
       + \frac{1}{2} \log|\mathbf{\Lambda}|
       \\
       &
       - \frac{1}{2} \sum^N_{n=1} \mathbf{x}_n^T \mathrm{diag}(\boldsymbol{\nu}) \mathbf{x}_n
       + \frac{1}{2} \sum^N_{n=1} \mathbf{x}_n^T \mathrm{diag}(\boldsymbol{\nu}) \mathbf{A} \mathbf{x}_{n-1}
       + \frac{1}{2} \sum^N_{n=1} \mathbf{x}_{n-1}^T\mathbf{A}^T \mathrm{diag}(\boldsymbol{\nu}) \mathbf{x}_n
       - \frac{1}{2} \sum^N_{n=1} \mathbf{x}_{n-1}^T\mathbf{A}^T \mathrm{diag}(\boldsymbol{\nu}) \mathbf{A} \mathbf{x}_{n-1}
       \\ &
       + \sum^N_{n=1} \sum^D_{d=1} \log\nu_d - \frac{1}{2} (N+1) D \log(2\pi)
       \\
       =&
       \begin{bmatrix}
         \mathbf{x}_0 \\ \mathbf{x}_1 \\ \vdots \\ \mathbf{x}_{N-1} \\ \mathbf{x}_N
       \end{bmatrix}^T
       \begin{bmatrix}
         -\frac{1}{2}\mathbf{\Lambda} - \frac{1}{2}\mathbf{A}\mathrm{diag}(\boldsymbol{\nu})\mathbf{A}^T
         &
         \frac{1}{2} \mathbf{A}^T\mathrm{diag}(\boldsymbol{\nu})
         &
         &
         &
         \\
         \frac{1}{2} \mathrm{diag}(\boldsymbol{\nu}) \mathbf{A}
         &
         -\frac{1}{2} \mathrm{diag}(\boldsymbol{\nu})
         - \frac{1}{2}\mathbf{A}^T\mathrm{diag}(\boldsymbol{\nu})\mathbf{A}^T
         &
         \frac{1}{2} \mathbf{A}^T\mathrm{diag}(\boldsymbol{\nu})
         &
         &
         \\
         &
         \ddots
         &
         \ddots
         &
         \ddots
         &
         \\
         &
         &
         \frac{1}{2} \mathrm{diag}(\boldsymbol{\nu}) \mathbf{A}
         &
         -\frac{1}{2} \mathrm{diag}(\boldsymbol{\nu})
         - \frac{1}{2}\mathbf{A}^T\mathrm{diag}(\boldsymbol{\nu})\mathbf{A}^T
         &
         \frac{1}{2} \mathbf{A}^T\mathrm{diag}(\boldsymbol{\nu})
         \\
         &
         &
         &
         \frac{1}{2} \mathrm{diag}(\boldsymbol{\nu}) \mathbf{A}
         &
         -\frac{1}{2} \mathrm{diag}(\boldsymbol{\nu})
       \end{bmatrix}
       \begin{bmatrix}
         \mathbf{x}_0 \\ \mathbf{x}_1 \\ \vdots \\ \mathbf{x}_{N-1} \\ \mathbf{x}_N
       \end{bmatrix}
       \\
       &
       + \frac{1}{2} \mathbf{x}_0^T \mathbf{\Lambda} \boldsymbol{\mu}
       + \frac{1}{2} \boldsymbol{\mu}^T \mathbf{\Lambda} \mathbf{x}_0
       - \frac{1}{2} \boldsymbol{\mu}^T \mathbf{\Lambda} \boldsymbol{\mu}
       + \frac{1}{2} \log|\mathbf{\Lambda}|
       + \sum^N_{n=1} \sum^D_{d=1} \log\nu_d - \frac{1}{2} (N+1) D \log(2\pi)


    For simplicity, :math:`\boldsymbol{\nu}` and :math:`\mathbf{A}` are assumed
    not to depend on :math:`n` in the above equation, but this distribution
    class supports that dependency. One only needs to do the following
    replacements in the equations: :math:`\boldsymbol{\nu} \leftarrow \boldsymbol{\nu}_n` and
    :math:`\mathbf{A} \leftarrow \mathbf{A}_n`, where :math:`n=1,\ldots,N`.

    .. math::

       u(\mathbf{X}) &=
       \begin{bmatrix}
         \begin{bmatrix} \mathbf{x}_0 & \ldots & \mathbf{x}_N \end{bmatrix}
         \\
         \begin{bmatrix} \mathbf{x}_0\mathbf{x}_0^T & \ldots & \mathbf{x}_N\mathbf{x}_N^T \end{bmatrix}
         \\
         \begin{bmatrix} \mathbf{x}_0\mathbf{x}_1^T & \ldots & \mathbf{x}_{N-1}\mathbf{x}_N^T \end{bmatrix}
       \end{bmatrix}
       \\
       \phi(\boldsymbol{\mu}, \mathbf{\Lambda}, \mathbf{A}, \boldsymbol{\nu}) &=
       \begin{bmatrix}
         \begin{bmatrix}
           \mathbf{\Lambda} \boldsymbol{\mu} & \mathbf{0} & \ldots & \mathbf{0}
         \end{bmatrix}
         \\
         \begin{bmatrix}
           -\frac{1}{2}\mathbf{\Lambda} - \frac{1}{2} \mathbf{A}\mathrm{diag}(\boldsymbol{\nu})\mathbf{A}^T &
           -\frac{1}{2}\mathrm{diag}(\boldsymbol{\nu}) - \frac{1}{2} \mathbf{A}\mathrm{diag}(\boldsymbol{\nu})\mathbf{A}^T &
           \ldots &
           -\frac{1}{2}\mathrm{diag}(\boldsymbol{\nu}) - \frac{1}{2} \mathbf{A}\mathrm{diag}(\boldsymbol{\nu})\mathbf{A}^T &
           -\frac{1}{2}\mathrm{diag}(\boldsymbol{\nu})
         \end{bmatrix}
         \\
         \begin{bmatrix}
           \mathbf{A}^T \mathrm{diag}(\boldsymbol{\nu}) & \ldots & \mathbf{A}^T \mathrm{diag}(\boldsymbol{\nu})
         \end{bmatrix}
       \end{bmatrix}
       \\
       g(\boldsymbol{\mu}, \mathbf{\Lambda}, \mathbf{A}, \boldsymbol{\nu}) &=
       \frac{1}{2}\log|\mathbf{\Lambda}| + \frac{1}{2} \sum^N_{n=1}\sum^D_{d=1}\log\nu_d
       \\
       f(\mathbf{X}) &= -\frac{1}{2} (N+1) D \log(2\pi)

    The log probability denisty function of the posterior approximation:

    .. math::

       \log q(\mathbf{X}) &=
       \begin{bmatrix}
         \mathbf{x}_0
         \\
         \mathbf{x}_1
         \\
         \vdots
         \\
         \mathbf{x}_{N-1}
         \\
         \mathbf{x}_N
       \end{bmatrix}^T
       \begin{bmatrix}
         \mathbf{\Phi}_0^{(2)} & \frac{1}{2}\mathbf{\Phi}_1^{(3)} & & &
         \\
         \frac{1}{2}{\mathbf{\Phi}_1^{(3)}}^T & \mathbf{\Phi}_1^{(2)} & \frac{1}{2}\mathbf{\Phi}_2^{(3)} & &
         \\
         & \ddots & \ddots & \ddots &
         \\
         & & \frac{1}{2}{\mathbf{\Phi}_{N-1}^{(3)}}^T & \mathbf{\Phi}_{N-1}^{(2)} & \frac{1}{2}\mathbf{\Phi}_N^{(3)}
         \\
         & & & \frac{1}{2}{\mathbf{\Phi}_N^{(3)}}^T & \mathbf{\Phi}_N^{(2)}
       \end{bmatrix}
       \begin{bmatrix}
         \mathbf{x}_0
         \\
         \mathbf{x}_1
         \\
         \vdots
         \\
         \mathbf{x}_{N-1}
         \\
         \mathbf{x}_N
       \end{bmatrix}
       + \ldots

    """


    def compute_message_to_parent(self, parent, index, u, u_mu_Lambda, u_A_nu, *u_inputs):
        r"""
        Compute a message to a parent.

        Parameters
        ----------
        index : int
            Index of the parent requesting the message.
        u : list of ndarrays
            Moments of this node.
        u_mu_Lambda : list of ndarrays
            Moments of parents :math:`(\boldsymbol{\mu}, \mathbf{\Lambda})`.
        u_A_nu : list of ndarrays
            Moments of parents :math:`(\mathbf{A}, \boldsymbol{\nu})`.
        u_inputs : list of ndarrays
            Moments of input signals.
        """

        D = np.shape(u[0])[-1]

        if index == 0:   # (mu, Lambda) -- GaussianWishartMoments
            x0 = u[0][...,0,:]
            x0x0 = u[1][...,0,:,:]
            m0 = x0
            m1 = -0.5
            m2 = -0.5 * x0x0
            m3 = 0.5
            return [m0, m1, m2, m3]

        elif index == 1: # (A, nu) -- GaussianGammaMoments
            XnXn = u[1]
            XpXn = u[2]

            # (..., N-1, D, D)
            m0 = XpXn.swapaxes(-1,-2)
            # (..., N-1, D, D, D)
            m1 = -0.5 * XnXn[..., :-1, None, :, :]
            # (..., N-1, D)
            m2 = -0.5 * np.einsum('...ii->...i', XnXn[...,1:,:,:])
            # (..., N-1, D)
            m3 = 0.5

            if len(u_inputs):
                Xn = u[0]
                z = u_inputs[0][0]
                zz = u_inputs[0][1]
                D_inputs = np.shape(z)[-1]
                m0_B = Xn[...,1:,:,None] * z[...,None,:]
                m1_BB = -0.5 * zz[..., None, :, :]
                m1_AB = -0.5 * Xn[..., :-1, None, :, None] * z[..., None, None, :]

                # Construct full message arrays from blocks
                m0 = np.concatenate([m0, m0_B], axis=-1)
                row1 = np.concatenate([m1, m1_AB], axis=-1)
                row2 = np.concatenate([m1_AB.swapaxes(-1,-2), m1_BB], axis=-1)
                m1 = np.concatenate([row1, row2], axis=-2)

            return [m0, m1, m2, m3]


        #     m1 = 0.5
        elif index == 2: # input signals

            # (..., N-1, D)
            Xn = u[0][...,1:,:]
            # (..., N-1, D)
            Xp = u[0][...,:-1,:]
            # (..., N-1, D, K)
            B = u_A_nu[0][...,D:]
            # (..., N-1, D, D, K)
            AB = u_A_nu[1][...,:D,D:]
            # (..., N-1, D, K, K)
            BB = u_A_nu[1][...,D:,D:]

            # (..., N-1, K)
            m0 = (
                np.einsum('...dk,...d->...k', B, Xn) -
                np.einsum('...dk,...d->...k', np.sum(AB, axis=-3), Xp)
            )
            # (..., N-1, K, K)
            m1 = -0.5 * np.sum(BB, axis=-3)

            return [m0, m1]

        raise IndexError("Parent index out of bounds")


    def compute_weights_to_parent(self, index, weights):

        if index == 0:   # mu_Lambda
            return weights
        elif index == 1: # A_nu
            return weights[...,np.newaxis,np.newaxis]
        elif index == 2: # input signals
            return weights[...,np.newaxis]
        else:
            raise ValueError("Index out of bounds")


    def compute_phi_from_parents(self, u_mu_Lambda, u_A_nu, *u_inputs, mask=True):
        """
        Compute the natural parameters using parents' moments.

        Parameters
        ----------
        u_parents : list of list of arrays
           List of parents' lists of moments.

        Returns
        -------
        phi : list of arrays
           Natural parameters.
        dims : tuple
           Shape of the variable part of phi.

        """

        # Dimensionality of the Gaussian states
        D = np.shape(u_mu_Lambda[0])[-1]

        # Number of time instances in the process
        N = self.N

        # Helpful variables (show shapes in comments)
        Lambda_mu = u_mu_Lambda[0]   # (..., D)
        Lambda = u_mu_Lambda[2]      # (..., D, D)
        nu_A = u_A_nu[0][...,:D]     # (..., N-1, D, D)
        nu_AA = u_A_nu[1][...,:D,:D] # (..., N-1, D, D, D)
        nu_B = u_A_nu[0][...,D:]     # (..., N-1, D, inputs)
        nu_BB = u_A_nu[1][...,D:,D:] # (..., N-1, D, inputs, inputs)
        nu_AB = u_A_nu[1][...,:D,D:] # (..., N-1, D, D, inputs)
        nu = u_A_nu[2] * np.ones(D)  # (..., N-1, D)
        # mu = u_mu[0]           # (..., D)
        # Lambda = u_Lambda[0]   # (..., D, D)
        # A = u_A[0][...,:D]     # (..., N-1, D, D)
        # AA = u_A[1][...,:D,:D] # (..., N-1, D, D, D)
        # B = u_A[0][...,D:]     # (..., N-1, D, inputs)
        # BB = u_A[1][...,D:,D:] # (..., N-1, D, inputs, inputs)
        # AB = u_A[1][...,:D,D:] # (..., N-1, D, D, inputs)
        # v = u_v[0]             # (..., N-1, D)
        if len(u_inputs):
            inputs = u_inputs[0][0]
        else:
            inputs = None

        # Allocate memory (take into account effective plates)
        if inputs is not None:
            plates_phi0 = misc.broadcasted_shape(np.shape(Lambda_mu)[:-1],
                                                 np.shape(nu_B)[:-3],
                                                 np.shape(nu_AB)[:-4])
        else:
            plates_phi0 = misc.broadcasted_shape(np.shape(Lambda_mu)[:-1])

        plates_phi1 = misc.broadcasted_shape(np.shape(Lambda)[:-2],
                                             np.shape(nu_AA)[:-4])
        plates_phi2 = misc.broadcasted_shape(np.shape(nu_A)[:-3])

        phi0 = np.zeros(plates_phi0+(N,D))
        phi1 = np.zeros(plates_phi1+(N,D,D))
        phi2 = np.zeros(plates_phi2+(N-1,D,D))

        # Parameters for x0
        phi0[...,0,:] = Lambda_mu #np.einsum('...ik,...k->...i', Lambda, mu)
        phi1[...,0,:,:] = -0.5 * Lambda

        # Effect of the input signals
        if inputs is not None:
            phi0[...,1:,:] += np.einsum('...ij,...j->...i', nu_B, inputs)
            phi0[...,:-1,:] -= np.einsum(
                '...ij,...j->...i',
                np.sum(nu_AB, axis=-3),
                inputs
            )

        # Diagonal blocks: -0.5 * (V_i + A_{i+1}' * V_{i+1} * A_{i+1})
        phi1[..., 1:, :, :] = -0.5 * misc.diag(nu, ndim=1)
        phi1[..., :-1, :, :] += -0.5 * np.sum(nu_AA, axis=-3) #np.einsum('...kij,...k->...ij', AA, v)
        #phi1 *= -0.5

        # Super-diagonal blocks: 0.5 * A.T * V
        # However, don't multiply by 0.5 because there are both super- and
        # sub-diagonal blocks (sum them together)
        phi2[..., :, :, :] = linalg.transpose(nu_A, ndim=1) # np.einsum('...ji,...j->...ij', A, v)

        return (phi0, phi1, phi2)

    def compute_cgf_from_parents(self, u_mu_Lambda, u_A_nu, *u_inputs):
        """
        Compute CGF using the moments of the parents.
        """
        g = _compute_cgf_for_gaussian_markov_chain(u_mu_Lambda[1],
                                                   u_mu_Lambda[3],
                                                   u_A_nu[3],
                                                   self.N)

        if len(u_inputs):
            D = np.shape(u_mu_Lambda[0])[-1]
            uu = u_inputs[0][1]
            nu_BB = u_A_nu[1][...,D:,D:]
            nu = u_A_nu[2]
            #BB_v = np.einsum('...d,...dij->...ij', v, BB)
            g_inputs = -0.5 * np.einsum(
                '...ij,...ij->...',
                uu,
                np.sum(nu_BB, axis=-3)
                #BB_v
            )
            # Sum over time axis
            if np.ndim(g_inputs) == 0 or np.shape(g_inputs)[-1] == 1:
                g_inputs *= self.N - 1
            if np.ndim(g_inputs) > 0:
                g_inputs = np.sum(g_inputs, axis=-1)
            g = g + g_inputs

        return g


    def plates_to_parent(self, index, plates):
        """
        Computes the plates of this node with respect to a parent.

        If this node has plates (...), the latent dimensionality is D
        and the number of time instances is N, the plates with respect
        to the parents are:
          (mu, Lambda): (...)
          (A, nu):      (...,N-1,D)

        Parameters
        ----------
        index : int
            The index of the parent node to use.
        """

        if index == 0:   # (mu, Lambda)
            return plates
        elif index == 1: # (A, nu)
            return plates + (self.N-1, self.D)
        elif index == 2: # input signals
            return plates + (self.N-1,)
        else:
            raise ValueError("Invalid parent index.")

    def plates_from_parent(self, index, plates):
        """
        Compute the plates using information of a parent node.

        If the plates of the parents are:
          (mu, Lambda): (...)
          (A, nu):      (...,N-1,D)
        the resulting plates of this node are (...)

        Parameters
        ----------
        index : int
            Index of the parent to use.
        """
        if index == 0:   # (mu, Lambda)
            return plates
        elif index == 1: # (A, nu)
            return plates[:-2]
        elif index == 2: # input signals
            return plates[:-1]
        else:
            raise ValueError("Invalid parent index.")


class GaussianMarkovChain(_TemplateGaussianMarkovChain):
    r"""
    Node for Gaussian Markov chain random variables.

    In a simple case, the graphical model can be presented as:

    .. bayesnet::

       \tikzstyle{latent} += [minimum size=30pt];

       \node[latent] (x0) {$\mathbf{x}_0$};
       \node[latent, right=of x0] (x1) {$\mathbf{x}_1$};
       \node[right=of x1] (dots) {$\cdots$};
       \node[latent, right=of dots] (xn) {$\mathbf{x}_{N-1}$};
       \edge {x0}{x1};
       \edge {x1}{dots};
       \edge {dots}{xn};

       \node[latent, above left=1 and 0.1 of x0] (mu) {$\boldsymbol{\mu}$};
       \node[latent, above right=1 and 0.1 of x0] (Lambda) {$\mathbf{\Lambda}$};
       \node[latent, above left=1 and 0.1 of dots] (A) {$\mathbf{A}$};
       \node[latent, above right=1 and 0.1 of dots] (nu) {$\boldsymbol{\nu}$};
       \edge {mu,Lambda} {x0};
       \edge {A,nu} {x1,dots,xn};

    where :math:`\boldsymbol{\mu}` and :math:`\mathbf{\Lambda}` are the mean and
    the precision matrix of the initial state, :math:`\mathbf{A}` is the state
    dynamics matrix and :math:`\boldsymbol{\nu}` is the precision of the
    innovation noise.  It is possible that :math:`\mathbf{A}` and/or
    :math:`\boldsymbol{\nu}` are different for each transition instead of being
    constant.

    The probability distribution is

    .. math::

       p(\mathbf{x}_0, \ldots, \mathbf{x}_{N-1}) = p(\mathbf{x}_0)
       \prod^{N-1}_{n=1} p(\mathbf{x}_n | \mathbf{x}_{n-1})

    where

    .. math::

       p(\mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_0 | \boldsymbol{\mu}, \mathbf{\Lambda})
       \\
       p(\mathbf{x}_n|\mathbf{x}_{n-1}) &= \mathcal{N}(\mathbf{x}_n |
       \mathbf{A}_{n-1}\mathbf{x}_{n-1}, \mathrm{diag}(\boldsymbol{\nu}_{n-1})).

    Parameters
    ----------

    mu : Gaussian-like node or (...,D)-array
        :math:`\boldsymbol{\mu}`, mean of :math:`x_0`, :math:`D`-dimensional
        with plates (...)

    Lambda : Wishart-like node or (...,D,D)-array
        :math:`\mathbf{\Lambda}`, precision matrix of :math:`x_0`,
        :math:`D\times D` -dimensional with plates (...)

    A : Gaussian-like node or (D,D)-array or (...,1,D,D)-array or (...,N-1,D,D)-array
        :math:`\mathbf{A}`, state dynamics matrix, :math:`D`-dimensional with
        plates (D,) or (...,1,D) or (...,N-1,D)

    nu : gamma-like node or (D,)-array or (...,1,D)-array or (...,N-1,D)-array
        :math:`\boldsymbol{\nu}`, diagonal elements of the precision of the
        innovation process, plates (D,) or (...,1,D) or (...,N-1,D)

    n : int, optional
        :math:`N`, the length of the chain. Must be given if :math:`\mathbf{A}`
        and :math:`\boldsymbol{\nu}` are constant over time.

    See also
    --------

    Gaussian, GaussianARD, Wishart, Gamma, SwitchingGaussianMarkovChain,
    VaryingGaussianMarkovChain, CategoricalMarkovChain
    """


    def __init__(self, mu, Lambda, A, nu, n=None, inputs=None, **kwargs):
        """
        Create GaussianMarkovChain node.
        """
        super().__init__(mu, Lambda, A, nu, n=n, inputs=inputs, **kwargs)


    @classmethod
    def _constructor(cls, mu, Lambda, A, nu, n=None, inputs=None, **kwargs):
        """
        Constructs distribution and moments objects.

        Compute the dimensions of phi and u.

        The plates and dimensions of the parents should be:
        mu:     (...)                    and D-dimensional
        Lambda: (...)                    and D-dimensional
        A:      (...,1,D) or (...,N-1,D) and D-dimensional
        v:      (...,1,D) or (...,N-1,D) and 0-dimensional
        N:      ()                       and 0-dimensional (dummy parent)

        Check that the dimensionalities of the parents are proper.
        For instance, A should be a collection of DxD matrices, thus
        the dimensionality and the last plate should both equal D.
        Similarly, `v` should be a collection of diagonal innovation
        matrix elements, thus the last plate should equal D.
        """

        mu_Lambda = WrapToGaussianWishart(mu, Lambda)
        A_nu = WrapToGaussianGamma(A, nu, ndim=1)

        D = mu_Lambda.dims[0][0]

        if inputs is not None:
            inputs = cls._ensure_moments(inputs, GaussianMoments, ndim=1)

        # Check whether to use input signals or not
        if inputs is None:
            _parent_moments = (GaussianWishartMoments((D,)),
                               GaussianGammaMoments((D,)))
        else:
            K = inputs.dims[0][0]
            _parent_moments = (GaussianWishartMoments((D,)),
                               GaussianGammaMoments((D,)),
                               GaussianMoments((K,)))

        # Time instances from input signals
        if inputs is not None and len(inputs.plates) >= 1:
            n_inputs = inputs.plates[-1]
        else:
            n_inputs = 1
        # Time instances from state dynamics matrix
        if len(A_nu.plates) >= 2:
            n_A_nu = A_nu.plates[-2]
        else:
            n_A_nu = 1
        # Check consistency of the number of time instances
        if n_inputs != n_A_nu and n_inputs != 1 and n_A_nu != 1:
            raise Exception("Plates of parents are giving different number of time instances")
        n_parents = max(n_A_nu, n_inputs)
        if n is None:
            if n_parents == 1:
                raise Exception("The number of time instances could not be "
                                "determined automatically. Give the number of "
                                "time instances.")
            n = n_parents + 1
        elif n_parents != 1 and n_parents+1 != n:
            raise Exception("The number of time instances must match "
                            "the number of last plates of parents: "
                            "%d != %d+1" % (n, n_parents))

        # Dimensionality of the states
        D = mu_Lambda.dims[0][0]
        # Number of states
        M = n
        # Dimensionality of the inputs
        if inputs is None:
            D_inputs = 0
        else:
            D_inputs = inputs.dims[0][0]

        # Check (mu, Lambda)
        if mu_Lambda.dims != ( (D,), (), (D, D), () ):
            raise Exception("Initial state parameters have wrong dimensionality")
        # Check (A, nu)
        if A_nu.dims != ( (D+D_inputs,), (D+D_inputs,D+D_inputs), (), () ):
            raise Exception("Dynamics matrix has wrong dimensionality")
        if len(A_nu.plates) == 0 or A_nu.plates[-1] != D:
            raise Exception("Dynamics matrix should have a last plate "
                            "equal to the dimensionality of the "
                            "system.")
        if (len(A_nu.plates) >= 2
            and A_nu.plates[-2] != 1
            and A_nu.plates[-2] != M-1):
            raise ValueError("The second last plate of the dynamics matrix "
                             "should have length equal to one or "
                             "N-1, where N is the number of time "
                             "instances.")
        # Check input signals
        if inputs is not None:
            if inputs.dims != ( (D_inputs,), (D_inputs, D_inputs) ):
                raise ValueError("Input signals have wrong dimensionality")

        moments = GaussianMarkovChainMoments(M, D)
        dims = ( (M,D), (M,D,D), (M-1,D,D) )
        distribution = GaussianMarkovChainDistribution(M, D)

        if inputs is None:
            parents = [mu_Lambda, A_nu]
        else:
            parents = [mu_Lambda, A_nu, inputs]

        return ( parents,
                 kwargs,
                 dims,
                 cls._total_plates(kwargs.get('plates'),
                                   distribution.plates_from_parent(0, mu_Lambda.plates),
                                   distribution.plates_from_parent(1, A_nu.plates)),
                 distribution,
                 moments,
                 _parent_moments)


    def rotate(self, R, inv=None, logdet=None):

        # It would be more efficient and simpler, if you just rotated the
        # moments and didn't touch phi. However, then you would need to call
        # update() before lower_bound_contribution. This is more error-safe.
        (u, phi, dg) = self._distribution.rotate(
            self.u,
            self.phi,
            R,
            inv=inv,
            logdet=logdet
        )
        self.u = u
        self.phi = phi
        self.g = self.g + dg

        return


class VaryingGaussianMarkovChainDistribution(TemplateGaussianMarkovChainDistribution):
    """
    Sub-classes implement distribution specific computations.
    """


    def compute_message_to_parent(self, parent, index, u, u_mu, u_Lambda, u_B,
                                   u_S, u_v):
        """
        Compute a message to a parent.

        Parameters
        -----------
        index : int
            Index of the parent requesting the message.
        u : list of ndarrays
            Moments of this node.
        u_mu : list of ndarrays
            Moments of parent `mu`.
        u_Lambda : list of ndarrays
            Moments of parent `Lambda`.
        u_B : list of ndarrays
            Moments of parent `B`.
        u_S : list of ndarrays
            Moments of parent `S`.
        u_v : list of ndarrays
            Moments of parent `v`.
        """

        if index == 0:   # mu
            raise NotImplementedError()
        elif index == 1: # Lambda
            raise NotImplementedError()
        elif index == 2: # B, (...,D)x(D,K)
            XnXn = u[1] # (...,N,D,D)
            XpXn = u[2] # (...,N,D,D)
            S = misc.atleast_nd(u_S[0], 2)  # (...,N,K)
            SS = misc.atleast_nd(u_S[1], 3) # (...,N,K,K)
            v = misc.atleast_nd(u_v[0], 2)  # (...,N,D)

            # m0: (...,D,D,K)
            m0 = np.einsum('...nji,...nk,...ni->...ijk',
                           XpXn,
                           S,
                           v)

            # m1: (...,D,D,K,D,K)

            if np.ndim(v) >= 2 and np.shape(v)[-2] > 1:
                raise ValueError("Innovation noise is time dependent")

            m1 = np.einsum('...nij,...nkl->...ikjl',
                           XnXn[...,:-1,:,:],
                           SS)
            m1 = -0.5 * np.einsum('...ikjl,...d->...dikjl',
                                  m1,
                                  v[...,0,:])

        elif index == 3: # S, (...,N-1)x(K)
            XnXn = u[1] # (...,N,D,D)
            XpXn = u[2] # (...,N,D,D)
            B = u_B[0]  # (...,D,D,K)
            BB = u_B[1] # (...,D,D,K,D,K)
            v = u_v[0]  # (...,N,D)

            # m0: (...,N,K)
            m0 = np.einsum('...nji,...ijk,...ni->...nk',
                           XpXn,
                           B,
                           np.atleast_2d(v))

            # m1: (...,N,K,K)

            if np.ndim(v) >= 2 and np.shape(v)[-2] > 1:
                raise ValueError("Innovation noise is time dependent")
            m1 = np.einsum('...dikjl,...d->...ikjl',
                           BB,
                           np.atleast_2d(v)[...,0,:])
            m1 = -0.5 * np.einsum('...nij,...ikjl->...nkl',
                                  XnXn[...,:-1,:,:],
                                  m1)

        elif index == 4: # v
            raise NotImplementedError()
        elif index == 5: # N
            raise NotImplementedError()

        return [m0, m1]


    def compute_weights_to_parent(self, index, weights):

        if index == 0: # mu
            return weights
        elif index == 1: # Lambda
            return weights
        elif index == 2: # B
            return weights[...,np.newaxis] # new plate axis for D
        elif index == 3: # S
            return weights[...,np.newaxis] # new plate axis for N
        elif index == 4: # v
            return weights[...,np.newaxis,np.newaxis] # new plate axis for N and D
        elif index == 5: # N
            return weights
        else:
            raise ValueError("Invalid index")


    def compute_phi_from_parents(self, u_mu, u_Lambda, u_B, u_S, u_v,
                                 mask=True):
        """
        Compute the natural parameters using parents' moments.

        Parameters
        ----------
        u_parents : list of list of arrays
           List of parents' lists of moments.

        Returns
        -------
        phi : list of arrays
           Natural parameters.
        dims : tuple
           Shape of the variable part of phi.

        """

        # Dimensionality of the Gaussian states
        D = np.shape(u_mu[0])[-1]

        # Number of time instances in the process
        N = self.N

        # Helpful variables (show shapes in comments)
        mu = u_mu[0]         # (..., D)
        Lambda = u_Lambda[0] # (..., D, D)
        B = u_B[0]           # (..., D, D, K)
        BB = u_B[1]          # (..., D, D, K, D, K)
        S = u_S[0]           # (..., N-1, K) or (..., 1, K)
        SS = u_S[1]          # (..., N-1, K, K)
        v = u_v[0]           # (..., N-1, D) or (..., 1, D)

        # TODO/FIXME: Take into account plates!
        plates_phi0 = misc.broadcasted_shape(np.shape(mu)[:-1],
                                             np.shape(Lambda)[:-2])
        plates_phi1 = misc.broadcasted_shape(np.shape(Lambda)[:-2],
                                             np.shape(v)[:-2],
                                             np.shape(BB)[:-5],
                                             np.shape(SS)[:-3])
        plates_phi2 = misc.broadcasted_shape(np.shape(B)[:-3],
                                             np.shape(S)[:-2],
                                             np.shape(v)[:-2])
        phi0 = np.zeros(plates_phi0 + (N,D))
        phi1 = np.zeros(plates_phi1 + (N,D,D))
        phi2 = np.zeros(plates_phi2 + (N-1,D,D))

        # Parameters for x0
        phi0[...,0,:] = np.einsum('...ik,...k->...i', Lambda, mu)
        phi1[...,0,:,:] = Lambda


        # Diagonal blocks: -0.5 * (V_i + A_{i+1}' * V_{i+1} * A_{i+1})
        phi1[..., 1:, :, :] = v[...,np.newaxis]*np.identity(D)
        if np.ndim(v) >= 2 and np.shape(v)[-2] > 1:
            raise Exception("This implementation is not efficient if "
                            "innovation noise is time-dependent.")
            phi1[..., :-1, :, :] += np.einsum('...dikjl,...kl,...d->...ij',
                                              BB[...,None,:,:,:,:,:],
                                              SS,
                                              v)
        else:
            # We know that S does not have the D plate so we can sum that plate
            # axis out
            v_BB = np.einsum('...dikjl,...d->...ikjl',
                             BB[...,None,:,:,:,:,:],
                             v)
            phi1[..., :-1, :, :] += np.einsum('...ikjl,...kl->...ij',
                                              v_BB,
                                              SS)

        #phi1[..., :-1, :, :] += np.einsum('...kij,...k->...ij', AA, v)
        phi1 *= -0.5

        # Super-diagonal blocks: 0.5 * A.T * V
        # However, don't multiply by 0.5 because there are both super- and
        # sub-diagonal blocks (sum them together)
        phi2[..., :, :, :] = np.einsum('...jik,...k,...j->...ij',
                                       B[...,None,:,:,:],
                                       S,
                                       v)
        #phi2[..., :, :, :] = np.einsum('...ji,...j->...ij', A, v)

        return (phi0, phi1, phi2)

    def compute_cgf_from_parents(self, u_mu, u_Lambda, u_B, u_S, u_v):
        """
        Compute CGF using the moments of the parents.
        """

        u_mumu_Lambda = linalg.inner(u_Lambda[0], u_mu[1], ndim=2)
        return _compute_cgf_for_gaussian_markov_chain(u_mumu_Lambda,
                                                      u_Lambda[1],
                                                      u_v[1],
                                                      self.N)

    def plates_to_parent(self, index, plates):
        """
        Computes the plates of this node with respect to a parent.

        If this node has plates (...), the latent dimensionality is D
        and the number of time instances is N, the plates with respect
        to the parents are:
          mu:     (...)
          Lambda: (...)
          A:      (...,N-1,D)
          v:      (...,N-1,D)

        Parameters
        -----------
        index : int
            The index of the parent node to use.
        """

        if index == 0:   # mu
            return plates
        elif index == 1: # Lambda
            return plates
        elif index == 2: # B
            return plates + (self.D,)
        elif index == 3: # S
            return plates + (self.N-1,)
        elif index == 4: # v
            return plates + (self.N-1,self.D)
        else:
            raise ValueError("Invalid parent index.")

    def plates_from_parent(self, index, plates):
        """
        Compute the plates using information of a parent node.

        If the plates of the parents are:
          mu:     (...)
          Lambda: (...)
          B:      (...,D)
          S:      (...,N-1)
          v:      (...,N-1,D)
          N:      ()
        the resulting plates of this node are (...)

        Parameters
        ----------
        index : int
            Index of the parent to use.
        """
        if index == 0:   # mu
            return plates
        elif index == 1: # Lambda
            return plates
        elif index == 2: # B, remove last plate D
            return plates[:-1]
        elif index == 3: # S, remove last plate N-1
            return plates[:-1]
        elif index == 4: # v, remove last plates N-1,D
            return plates[:-2]
        else:
            raise ValueError("Invalid parent index.")




class VaryingGaussianMarkovChain(_TemplateGaussianMarkovChain):
    r"""
    Node for Gaussian Markov chain random variables with time-varying dynamics.

    The node models a sequence of Gaussian variables
    :math:`\mathbf{x}_0,\ldots,\mathbf{x}_{N-1}` with linear Markovian dynamics.
    The time variability of the dynamics is obtained by modelling the state
    dynamics matrix as a linear combination of a set of matrices with
    time-varying linear combination weights.  The
    graphical model can be presented as:

    .. bayesnet::

       \tikzstyle{latent} += [minimum size=40pt];

       \node[latent] (x0) {$\mathbf{x}_0$};
       \node[latent, right=of x0] (x1) {$\mathbf{x}_1$};
       \node[right=of x1] (dots) {$\cdots$};
       \node[latent, right=of dots] (xn) {$\mathbf{x}_{N-1}$};
       \edge {x0}{x1};
       \edge {x1}{dots};
       \edge {dots}{xn};

       \node[latent, above left=1 and 0.1 of x0] (mu) {$\boldsymbol{\mu}$};
       \node[latent, above right=1 and 0.1 of x0] (Lambda) {$\mathbf{\Lambda}$};
       \node[det, below=of x1] (A0) {$\mathbf{A}_0$};
       \node[right=of A0] (Adots) {$\cdots$};
       \node[det, right=of Adots] (An) {$\mathbf{A}_{N-2}$};
       \node[latent, above=of dots] (nu) {$\boldsymbol{\nu}$};
       \edge {mu,Lambda} {x0};
       \edge {nu} {x1,dots,xn};
       \edge {A0} {x1};
       \edge {Adots} {dots};
       \edge {An} {xn};

       \node[latent, below=of A0] (s0) {$s_{0,k}$};
       \node[right=of s0] (sdots) {$\cdots$};
       \node[latent, right=of sdots] (sn) {$\mathbf{s}_{N-2,k}$};
       \node[latent, left=of s0] (B) {$\mathbf{B}_k$};
       \edge {B} {A0, Adots, An};
       \edge {s0} {A0};
       \edge {sdots} {Adots};
       \edge {sn} {An};

       \plate {K} {(B)(s0)(sdots)(sn)} {$k=0,\ldots,K-1$};

    where :math:`\boldsymbol{\mu}` and :math:`\mathbf{\Lambda}` are the mean and
    the precision matrix of the initial state, :math:`\boldsymbol{\nu}` is the
    precision of the innovation noise, and :math:`\mathbf{A}_n` are the state
    dynamics matrix obtained by mixing matrices :math:`\mathbf{B}_k` with
    weights :math:`s_{n,k}`.

    The probability distribution is

    .. math::

       p(\mathbf{x}_0, \ldots, \mathbf{x}_{N-1}) = p(\mathbf{x}_0)
       \prod^{N-1}_{n=1} p(\mathbf{x}_n | \mathbf{x}_{n-1})

    where

    .. math::

       p(\mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_0 | \boldsymbol{\mu}, \mathbf{\Lambda})
       \\
       p(\mathbf{x}_n|\mathbf{x}_{n-1}) &= \mathcal{N}(\mathbf{x}_n |
       \mathbf{A}_{n-1}\mathbf{x}_{n-1}, \mathrm{diag}(\boldsymbol{\nu})),
       \quad \text{for } n=1,\ldots,N-1,
       \\
       \mathbf{A}_n & = \sum^{K-1}_{k=0} s_{n,k} \mathbf{B}_k, \quad \text{for }
       n=0,\ldots,N-2.


    Parameters
    ----------

    mu : Gaussian-like node or (...,D)-array
        :math:`\boldsymbol{\mu}`, mean of :math:`x_0`, :math:`D`-dimensional
        with plates (...)

    Lambda : Wishart-like node or (...,D,D)-array
        :math:`\mathbf{\Lambda}`, precision matrix of :math:`x_0`,
        :math:`D\times D` -dimensional with plates (...)

    B : Gaussian-like node or (...,D,D,K)-array
        :math:`\{\mathbf{B}_k\}_{k=0}^{K-1}`, a set of state dynamics matrix,
        :math:`D \times K`-dimensional with plates (...,D)

    S : Gaussian-like node or (...,N-1,K)-array

        :math:`\{\mathbf{s}_0,\ldots,\mathbf{s}_{N-2}\}`, time-varying weights
        of the linear combination, :math:`K`-dimensional with plates (...,N-1)

    nu : gamma-like node or (...,D)-array
        :math:`\boldsymbol{\nu}`, diagonal elements of the precision of the
        innovation process, plates (...,D)

    n : int, optional
        :math:`N`, the length of the chain. Must be given if :math:`\mathbf{S}`
        does not have plates over the time domain (which would not make sense).

    See also
    --------

    Gaussian, GaussianARD, Wishart, Gamma, GaussianMarkovChain,
    SwitchingGaussianMarkovChain

    Notes
    -----

    Equivalent model block can be constructed with :class:`GaussianMarkovChain`
    by explicitly using :class:`SumMultiply` to compute the linear combination.
    However, that approach is not very efficient for large datasets because it
    does not utilize the structure of :math:`\mathbf{A}_n`, thus it explicitly
    computes huge moment arrays.

    References
    ----------

    :cite:`Luttinen:2014`

    """

    def __init__(self, mu, Lambda, B, S, nu, n=None, **kwargs):
        """
        Create VaryingGaussianMarkovChain node.
        """
        super().__init__(mu, Lambda, B, S, nu, n=n, **kwargs)


    @classmethod
    def _constructor(cls, mu, Lambda, B, S, v, n=None, **kwargs):
        """
        Constructs distribution and moments objects.

        Compute the dimensions of phi and u.

        The plates and dimensions of the parents should be:
        mu:     (...)                    and D-dimensional
        Lambda: (...)                    and D-dimensional
        B:      (...,D)                  and (D,K)-dimensional
        S:      (...,N-1)                and K-dimensional
        v:      (...,1,D) or (...,N-1,D) and 0-dimensional
        N:      ()                       and 0-dimensional (dummy parent)

        Check that the dimensionalities of the parents are proper.
        """

        mu = cls._ensure_moments(mu, GaussianMoments, ndim=1)
        Lambda = cls._ensure_moments(Lambda, WishartMoments, ndim=1)
        B = cls._ensure_moments(B, GaussianMoments, ndim=2)
        S = cls._ensure_moments(S, GaussianMoments, ndim=1)
        v = cls._ensure_moments(v, GammaMoments)

        (D, K) = B.dims[0]

        parent_moments = (
            GaussianMoments((D,)),
            WishartMoments((D,)),
            GaussianMoments((D, K)),
            GaussianMoments((K,)),
            GammaMoments()
        )

        # A dummy wrapper for the number of time instances.
        n_S = 1
        if len(S.plates) >= 1:
            n_S = S.plates[-1]
        n_v = 1
        if len(v.plates) >= 2:
            n_v = v.plates[-2]
        if n_v != n_S and n_v != 1 and n_S != 1:
            raise Exception(
                "Plates of A and v are giving different number of time "
                "instances")
        n_S = max(n_v, n_S)
        if n is None:
            if n_S == 1:
                raise Exception(
                    "The number of time instances could not be determined "
                    "automatically. Give the number of time instances.")

            n = n_S + 1
        elif n_S != 1 and n_S+1 != n:
            raise Exception(
                "The number of time instances must match the number of last "
                "plates of parents:" "%d != %d+1"
                % (n, n_S))

        D = mu.dims[0][0]
        K = B.dims[0][-1]
        M = n #N.get_moments()[0]

        # Check mu
        if mu.dims != ( (D,), (D,D) ):
            raise ValueError("First parent has wrong dimensionality")
        # Check Lambda
        if Lambda.dims != ( (D,D), () ):
            raise ValueError("Second parent has wrong dimensionality")
        # Check B
        if B.dims != ( (D,K), (D,K,D,K) ):
            raise ValueError("Third parent has wrong dimensionality {0}. Should be {1}.".format(B.dims[0], (D,K)))
        if len(B.plates) == 0 or B.plates[-1] != D:
            raise ValueError("Third parent should have a last plate "
                             "equal to the dimensionality of the "
                             "system.")
        if S.dims != ( (K,), (K,K) ):
            raise ValueError("Fourth parent has wrong dimensionality %s, "
                             "should be %s"
                             % (S.dims, ( (K,), (K,K) )))
        if (len(S.plates) >= 1
            and S.plates[-1] != 1
            and S.plates[-1] != M-1):
            raise ValueError("The last plate of the fourth "
                             "parent should have length equal to one or "
                             "N-1, where N is the number of time "
                             "instances.")
        # Check v
        if v.dims != ( (), () ):
            raise Exception("Fifth parent has wrong dimensionality")
        if len(v.plates) == 0 or v.plates[-1] != D:
            raise Exception("Fifth parent should have a last plate "
                            "equal to the dimensionality of the "
                            "system.")
        if (len(v.plates) >= 2
            and v.plates[-2] != 1
            and v.plates[-2] != M-1):
            raise ValueError("The second last plate of the fifth "
                             "parent should have length equal to one or "
                             "N-1 where N is the number of time "
                             "instances.")

        distribution = VaryingGaussianMarkovChainDistribution(M, D)
        moments = GaussianMarkovChainMoments(M, D)

        parents = [mu, Lambda, B, S, v]

        dims = ( (M,D), (M,D,D), (M-1,D,D) )

        return (parents,
                kwargs,
                dims,
                cls._total_plates(kwargs.get('plates'),
                                  distribution.plates_from_parent(0, mu.plates),
                                  distribution.plates_from_parent(1, Lambda.plates),
                                  distribution.plates_from_parent(2, B.plates),
                                  distribution.plates_from_parent(3, S.plates),
                                  distribution.plates_from_parent(4, v.plates)),
                distribution,
                moments,
                parent_moments)



class SwitchingGaussianMarkovChainDistribution(TemplateGaussianMarkovChainDistribution):
    """
    Sub-classes implement distribution specific computations.
    """


    def __init__(self, N, D, K):
        self.K = K
        super().__init__(N, D)

    def compute_message_to_parent(self, parent, index, u, u_mu, u_Lambda, u_B,
                                   u_Z, u_v):
        """
        Compute a message to a parent.

        Parameters
        ----------
        index : int
            Index of the parent requesting the message.
        u : list of ndarrays
            Moments of this node.
        u_mu : list of ndarrays
            Moments of parent `mu`.
        u_Lambda : list of ndarrays
            Moments of parent `Lambda`.
        u_B : list of ndarrays
            Moments of parent `B`.
        u_Z : list of ndarrays
            Moments of parent `Z`.
        u_v : list of ndarrays
            Moments of parent `v`.
        """

        if index == 0:   # mu
            raise NotImplementedError()
        elif index == 1: # Lambda
            raise NotImplementedError()

        elif index == 2: # B, (...,K,D)x(D)

            XnXn = u[1]                   # (...,N,D,D)
            XpXn = u[2]                   # (...,N-1,D,D)
            Z = u_Z[0]                    # (...,N-1,K)
            v = misc.atleast_nd(u_v[0], 2)  # (...,N-1,D)

            # Check that there is no time-dependency in v and remove the axis
            if np.ndim(v) >= 2 and np.shape(v)[-2] > 1:
                raise ValueError("Innovation noise is time dependent")
            v = np.squeeze(v, axis=-2)

            # m0: (...,K,D,D)
            m0 = np.einsum('...nji,...nk,...i->...kij',
                           XpXn,
                           Z,
                           v)

            # m1: (...,K,D,D,D)

            m1 = np.einsum('...nij,...nk->...kij',
                           XnXn[...,:-1,:,:],
                           Z)
            m1 = -0.5 * np.einsum('...kij,...d->...kdij',
                                  m1,
                                  v)

            return [m0, m1]

        elif index == 3: # Z, (...,N-1)x(K)

            XnXn = u[1]                     # (...,N,D,D)
            XpXn = u[2]                     # (...,N-1,D,D)
            B = u_B[0]                      # (...,K,D,D)
            BB = u_B[1]                     # (...,K,D,D,D)
            v = misc.atleast_nd(u_v[0], 2)    # (...,N-1,D)
            logv = misc.atleast_nd(u_v[1], 2) # (...,N-1,D)

            # Check that there is no time-dependency in v and remove the axis
            if np.ndim(v) >= 2 and np.shape(v)[-2] > 1:
                raise ValueError("Innovation noise is time dependent")
            v = np.squeeze(v, axis=-2)
            if np.ndim(logv) >= 2 and np.shape(logv)[-2] > 1:
                raise ValueError("Innovation noise is time dependent")
            logv = np.squeeze(logv, axis=-2)

            XnXn_v = np.einsum('...nii,...i->...n',
                               XnXn[...,1:,:,:],
                               v)
            XpXn_v_B = np.einsum('...nil,...l,...kli->...nk',
                                 XpXn,
                                 v,
                                 B)
            BvB = np.einsum('...kdij,...d->...kij',
                            BB,
                            v)
            XpXp_BvB = np.einsum('...nij,...kij->...nk',
                                 XnXn[...,:-1,:,:],
                                 BvB)

            m0 = ( -0.5 * XnXn_v[...,None]
                   + XpXn_v_B
                   -0.5 * XpXp_BvB
                   +0.5 * np.sum(logv, axis=-1)[...,None,None]
                   -0.5 * self.D * np.log(2*np.pi) )

            return [m0]

        elif index == 4: # v
            raise NotImplementedError()
        elif index == 5: # N
            raise NotImplementedError()



    def compute_weights_to_parent(self, index, weights):

        if index == 0: # mu: (...)x(N,D) -> (...)x(D)
            return weights
        elif index == 1: # Lambda: (...)x(N,D) -> (...)x(D,D)
            return weights
        elif index == 2: # B: (...)x(N,D) -> (...,K,D)x(D)
            return weights[...,None,None]
        elif index == 3: # Z: (...)x(N,D) -> (...,N-1)x(K)
            return weights[...,None]
        elif index == 4: # v: (...)x(N,D) -> (...,N-1,D)x()
            return weights[...,None,None]
        else:
            raise ValueError("Invalid index")


    def compute_phi_from_parents(self, u_mu, u_Lambda, u_B, u_Z, u_v,
                                 mask=True):
        """
        Compute the natural parameters using parents' moments.

        Parameters
        ----------
        u_parents : list of list of arrays
           List of parents' lists of moments.

        Returns
        -------
        phi : list of arrays
           Natural parameters.
        dims : tuple
           Shape of the variable part of phi.

        """

        # Dimensionality of the Gaussian states
        D = np.shape(u_mu[0])[-1]

        # Number of time instances in the process
        N = self.N

        # Helpful variables (show shapes in comments)
        mu = u_mu[0]                    # (..., D)
        Lambda = u_Lambda[0]            # (..., D, D)
        B = u_B[0]                      # (..., K, D, D)
        BB = u_B[1]                     # (..., K, D, D, D)
        Z = u_Z[0]                      # (..., N-1, K)
        v = misc.atleast_nd(u_v[0], 2) # (..., N-1, D) or (..., 1, D)

        # TODO/FIXME: Take into account plates!
        plates_phi0 = misc.broadcasted_shape(np.shape(mu)[:-1],
                                             np.shape(Lambda)[:-2])
        plates_phi1 = misc.broadcasted_shape(np.shape(Lambda)[:-2],
                                             np.shape(v)[:-2],
                                             np.shape(BB)[:-4],
                                             np.shape(Z)[:-2])
        plates_phi2 = misc.broadcasted_shape(np.shape(B)[:-3],
                                             np.shape(Z)[:-2],
                                             np.shape(v)[:-2])
        phi0 = np.zeros(plates_phi0 + (N,D))
        phi1 = np.zeros(plates_phi1 + (N,D,D))
        phi2 = np.zeros(plates_phi2 + (N-1,D,D))

        # Parameters for x0
        phi0[...,0,:] = np.einsum('...ik,...k->...i', Lambda, mu)
        phi1[...,0,:,:] = Lambda


        # Diagonal blocks: -0.5 * (V_i + A_{i+1}' * V_{i+1} * A_{i+1})
        phi1[..., 1:, :, :] = v[...,None]*np.identity(D)
        if np.shape(v)[-2] > 1:
            raise Exception("This implementation is not efficient if "
                            "innovation noise is time-dependent.")
            phi1[..., :-1, :, :] += np.einsum('...kdij,...nk,...nd->...nij',
                                              BB[...,:,:,:,:],
                                              Z,
                                              v)
        else:
            # We know that S does not have the D plate so we can sum that plate
            # axis out
            v_BB = np.einsum('...kdij,...nd->...nkij',
                             BB[...,:,:,:,:],
                             v)
            phi1[..., :-1, :, :] += np.einsum('...nkij,...nk->...nij',
                                              v_BB,
                                              Z)

        phi1 *= -0.5

        # Super-diagonal blocks: 0.5 * A.T * V
        # However, don't multiply by 0.5 because there are both super- and
        # sub-diagonal blocks (sum them together)
        phi2[..., :, :, :] = np.einsum('...kji,...nk,...nj->...nij',
                                       B[...,:,:,:],
                                       Z,
                                       v)

        return (phi0, phi1, phi2)

    def compute_cgf_from_parents(self, u_mu, u_Lambda, u_B, u_Z, u_v):
        """
        Compute CGF using the moments of the parents.
        """

        u_mumu_Lambda = linalg.inner(u_Lambda[0], u_mu[1], ndim=2)
        return _compute_cgf_for_gaussian_markov_chain(u_mumu_Lambda,
                                                      u_Lambda[1],
                                                      u_v[1],
                                                      self.N)

    def plates_to_parent(self, index, plates):
        """
        Computes the plates of this node with respect to a parent.

        If this node has plates (...), the latent dimensionality is D
        and the number of time instances is N, the plates with respect
        to the parents are:
          mu:     (...)
          Lambda: (...)
          A:      (...,N-1,D)
          v:      (...,N-1,D)

        Parameters
        ----------
        index : int
            The index of the parent node to use.
        """

        if index == 0:   # mu: (...)x(N,D) -> (...)x(D)
            return plates
        elif index == 1: # Lambda: (...)x(N,D) -> (...)x(D,D)
            return plates
        elif index == 2: # B: (...)x(N,D) -> (...,K,D)x(D)
            return plates + (self.K,self.D)
        elif index == 3: # Z: (...)x(N,D) -> (...,N-1)x(K)
            return plates + (self.N-1,)
        elif index == 4: # v: (...)x(N,D) -> (...,N-1,D)x()
            return plates + (self.N-1,self.D)
        else:
            raise ValueError("Invalid parent index.")


    def plates_from_parent(self, index, plates):
        """
        Compute the plates using information of a parent node.

        If the plates of the parents are:
          mu:     (...)
          Lambda: (...)
          B:      (...,D)
          S:      (...,N-1)
          v:      (...,N-1,D)
          N:      ()
        the resulting plates of this node are (...)

        Parameters
        ----------
        index : int
            Index of the parent to use.
        """
        if index == 0:   # mu: (...)x(D) -> (...)x(N,D)
            return plates
        elif index == 1: # Lambda: (...)x(D,D) -> (...)x(N,D)
            return plates
        elif index == 2: # B: (...,K,D)x(D) -> (...)x(N,D)
            return plates[:-2]
        elif index == 3: # Z: (...,N-1)x(K) -> (...)x(N,D)
            return plates[:-1]
        elif index == 4: # v: (...,N-1,D)x() -> (...)x(N,D)
            return plates[:-2]
        else:
            raise ValueError("Invalid parent index.")




class SwitchingGaussianMarkovChain(_TemplateGaussianMarkovChain):
    r"""
    Node for Gaussian Markov chain random variables with switching dynamics.

    The node models a sequence of Gaussian variables
    :math:`\mathbf{x}_0,\ldots,\mathbf{x}_{N-1}$ with linear Markovian dynamics.
    The dynamics may change in time, which is obtained by having a set of
    matrices and at each time selecting one of them as the state dynamics
    matrix.  The graphical model can be presented as:

    .. bayesnet::

       \tikzstyle{latent} += [minimum size=40pt];

       \node[latent] (x0) {$\mathbf{x}_0$};
       \node[latent, right=of x0] (x1) {$\mathbf{x}_1$};
       \node[right=of x1] (dots) {$\cdots$};
       \node[latent, right=of dots] (xn) {$\mathbf{x}_{N-1}$};
       \edge {x0}{x1};
       \edge {x1}{dots};
       \edge {dots}{xn};

       \node[latent, above left=1 and 0.1 of x0] (mu) {$\boldsymbol{\mu}$};
       \node[latent, above right=1 and 0.1 of x0] (Lambda) {$\mathbf{\Lambda}$};
       \node[det, below=of x1] (A0) {$\mathbf{A}_0$};
       \node[right=of A0] (Adots) {$\cdots$};
       \node[det, right=of Adots] (An) {$\mathbf{A}_{N-2}$};
       \node[latent, above=of dots] (nu) {$\boldsymbol{\nu}$};
       \edge {mu,Lambda} {x0};
       \edge {nu} {x1,dots,xn};
       \edge {A0} {x1};
       \edge {Adots} {dots};
       \edge {An} {xn};

       \node[latent, below=of A0] (z0) {$z_0$};
       \node[right=of z0] (zdots) {$\cdots$};
       \node[latent, right=of zdots] (zn) {$z_{N-2}$};
       \node[latent, left=of z0] (B) {$\mathbf{B}_k$};
       \edge {B} {A0, Adots, An};
       \edge {z0} {A0};
       \edge {zdots} {Adots};
       \edge {zn} {An};

       \plate {K} {(B)} {$k=0,\ldots,K-1$};

    where :math:`\boldsymbol{\mu}` and :math:`\mathbf{\Lambda}` are the mean and
    the precision matrix of the initial state, :math:`\boldsymbol{\nu}` is the
    precision of the innovation noise, and :math:`\mathbf{A}_n` are the state
    dynamics matrix obtained by selecting one of the matrices
    :math:`\{\mathbf{B}_k\}^{K-1}_{k=0}` at each time.  The selections are
    provided by :math:`z_n\in\{0,\ldots,K-1\}`.  The probability distribution is

    .. math::

       p(\mathbf{x}_0, \ldots, \mathbf{x}_{N-1}) = p(\mathbf{x}_0)
       \prod^{N-1}_{n=1} p(\mathbf{x}_n | \mathbf{x}_{n-1})

    where

    .. math::

       p(\mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_0 | \boldsymbol{\mu}, \mathbf{\Lambda})
       \\
       p(\mathbf{x}_n|\mathbf{x}_{n-1}) &= \mathcal{N}(\mathbf{x}_n |
       \mathbf{A}_{n-1}\mathbf{x}_{n-1}, \mathrm{diag}(\boldsymbol{\nu})),
       \quad \text{for } n=1,\ldots,N-1,
       \\
       \mathbf{A}_n &= \mathbf{B}_{z_n}, \quad \text{for }
       n=0,\ldots,N-2.


    Parameters
    ----------

    mu : Gaussian-like node or (...,D)-array
        :math:`\boldsymbol{\mu}`, mean of :math:`x_0`, :math:`D`-dimensional
        with plates (...)

    Lambda : Wishart-like node or (...,D,D)-array
        :math:`\mathbf{\Lambda}`, precision matrix of :math:`x_0`,
        :math:`D\times D` -dimensional with plates (...)

    B : Gaussian-like node or (...,D,D,K)-array
        :math:`\{\mathbf{B}_k\}_{k=0}^{K-1}`, a set of state dynamics matrix,
        :math:`D \times K`-dimensional with plates (...,D)

    Z : categorical-like node or (...,N-1)-array
        :math:`\{z_0,\ldots,z_{N-2}\}`, time-dependent selection,
        :math:`K`-categorical with plates (...,N-1)

    nu : gamma-like node or (...,D)-array
        :math:`\boldsymbol{\nu}`, diagonal elements of the precision of the
        innovation process, plates (...,D)

    n : int, optional
        :math:`N`, the length of the chain. Must be given if :math:`\mathbf{Z}`
        does not have plates over the time domain (which would not make sense).

    See also
    --------

    Gaussian, GaussianARD, Wishart, Gamma, GaussianMarkovChain,
    VaryingGaussianMarkovChain, Categorical, CategoricalMarkovChain

    Notes
    -----

    Equivalent model block can be constructed with :class:`GaussianMarkovChain`
    by explicitly using :class:`Gate` to select the state dynamics matrix.
    However, that approach is not very efficient for large datasets because it
    does not utilize the structure of :math:`\mathbf{A}_n`, thus it explicitly
    computes huge moment arrays.
    """


    def __init__(self, mu, Lambda, B, Z, nu, n=None, **kwargs):
        """
        Create SwitchingGaussianMarkovChain node.
        """
        super().__init__(mu, Lambda, B, Z, nu, n=n, **kwargs)


    @classmethod
    def _constructor(cls, mu, Lambda, B, Z, v, n=None, **kwargs):
        """
        Constructs distribution and moments objects.

        Compute the dimensions of phi and u.

        The plates and dimensions of the parents should be:
        mu:     (...)                    and D-dimensional
        Lambda: (...)                    and D-dimensional
        B:      (...,K,D)                and D-dimensional
        Z:      (...,N-1)                and K-dimensional
        v:      (...,1,D) or (...,N-1,D) and 0-dimensional

        Check that the dimensionalities of the parents are proper.
        """

        # Infer the number of dynamic matrices
        B = cls._ensure_moments(B, GaussianMoments, ndim=1)
        K = B.plates[-2]

        mu = cls._ensure_moments(mu, GaussianMoments, ndim=1)
        Lambda = cls._ensure_moments(Lambda, WishartMoments, ndim=1)
        Z = cls._ensure_moments(Z, CategoricalMoments, categories=K)
        v = cls._ensure_moments(v, GammaMoments)

        parent_moments = (
            mu._moments,
            Lambda._moments,
            B._moments,
            Z._moments,
            v._moments
        )

        # Infer the length of the chain
        n_Z = 1
        if len(Z.plates) == 0:
            raise ValueError("Z must have temporal axis on plates")
        n_Z = Z.plates[-1]
        n_v = 1
        if len(v.plates) >= 2:
            n_v = v.plates[-2]
        if n_v != n_Z and n_v != 1 and n_Z != 1:
            raise Exception(
                "Plates of Z and v are giving different number of time "
                "instances")
        n_Z = max(n_v, n_Z)
        if n is None:
            if n_Z == 1:
                raise Exception(
                    "The number of time instances could not be determined "
                    "automatically. Give the number of time instances.")

            n = n_Z + 1
        elif n_Z != 1 and n_Z+1 != n:
            raise Exception(
                "The number of time instances must match the number of last "
                "plates of parents:" "%d != %d+1"
                % (n, n_Z))


        D = mu.dims[0][0]
        K = Z.dims[0][0]
        M = n #N.get_moments()[0]

        # Check mu
        if mu.dims != ( (D,), (D,D) ):
            raise ValueError("First parent has wrong dimensionality")
        # Check Lambda
        if Lambda.dims != ( (D,D), () ):
            raise ValueError("Second parent has wrong dimensionality")
        # Check B
        if B.dims != ( (D,), (D,D) ):
            raise ValueError("Third parent has wrong dimensionality")
        if len(B.plates) < 2 or B.plates[-2:] != (K,D):
            raise ValueError("Third parent should have a last plate "
                             "equal to the dimensionality of the "
                             "system.")
        if Z.dims != ( (K,), ):
            raise ValueError("Fourth parent has wrong dimensionality %s, "
                             "should be %s"
                             % (Z.dims, ( (K,), )))
        if Z.plates[-1] != M-1:
            raise ValueError("The last plate of the fourth "
                             "parent should have length equal to one or "
                             "N-1, where N is the number of time "
                             "instances.")
        # Check v
        if v.dims != ( (), () ):
            raise Exception("Fifth parent has wrong dimensionality")
        if len(v.plates) == 0 or v.plates[-1] != D:
            raise Exception("Fifth parent should have a last plate "
                            "equal to the dimensionality of the "
                            "system.")
        if (len(v.plates) >= 2
            and v.plates[-2] != 1
            and v.plates[-2] != M-1):
            raise ValueError("The second last plate of the fifth "
                             "parent should have length equal to one or "
                             "N-1 where N is the number of time "
                             "instances.")


        dims = ( (M,D), (M,D,D), (M-1,D,D) )
        distribution = SwitchingGaussianMarkovChainDistribution(M, D, K)
        moments = GaussianMarkovChainMoments(M, D)

        parents = [mu, Lambda, B, Z, v]

        return (parents,
                kwargs,
                dims,
                cls._total_plates(kwargs.get('plates'),
                                  distribution.plates_from_parent(0, mu.plates),
                                  distribution.plates_from_parent(1, Lambda.plates),
                                  distribution.plates_from_parent(2, B.plates),
                                  distribution.plates_from_parent(3, Z.plates),
                                  distribution.plates_from_parent(4, v.plates)),
                distribution,
                moments,
                parent_moments)


class _MarkovChainToGaussian(Deterministic):
    """
    Transform a Gaussian Markov chain node into a Gaussian node.

    This node is deterministic.
    """

    def __init__(self, X, **kwargs):

        X = self._ensure_moments(X, GaussianMarkovChainMoments)

        D = X.dims[0][-1]

        self._moments = GaussianMoments((D,))
        self._parent_moments = (X._moments,)

        super().__init__(X, dims=self._moments.dims, **kwargs)


    def _plates_to_parent(self, index):
        """
        Return the number of plates to the parent node.

        Normally, the parent sees the same number of plates as the
        node itself.  However, now that one of the variable dimensions
        of the parents corresponds to a plate in this node, it is
        necessary to fix it here: the last plate is ignored when
        calculating plates with respect to the parent.

        Parent:
        Plates = (...)
        Dims = (N, ...)
        This node:
        Plates = (..., N)
        Dims = (...)
        """
        return self.plates[:-1]

    def _plates_from_parent(self, index):
        # Sub-classes may want to overwrite this if they manipulate plates
        if index != 0:
            raise ValueError("Invalid parent index.")

        parent = self.parents[0]
        plates = parent.plates + (parent.dims[0][0],)
        return plates

    def _compute_moments(self, u):
        """
        Transform the moments of a GMC to moments of a Gaussian.

        There is no need to worry about the plates and variable
        dimensions because the child node is free to interpret the
        axes as it pleases.  However, the Gaussian moments contain
        only <X(n)> and <X(n)*X(n)> but not <X(n-1)X(n)>, thus the
        last moment is discarded.
        """

        # Get the moments from the parent Gaussian Markov Chain
        #u = self.parents[0].get_moments() #message_to_child()

        # Send only moments <X(n)> and <X(n)X(n)> but not <X(n-1)X(n)>
        return u[:2]


    def _compute_weights_to_parent(self, index, weights):
        # Remove the last axis of the mask
        if np.ndim(weights) >= 1:
            weights = np.sum(weights, axis=-1)
        return weights


    @staticmethod
    def _compute_message_to_parent(index, m_children, *u_parents):
        """
        Transform a message to a Gaussian into a message to a GMC.

        The messages to a Gaussian are almost correct, there are only
        two minor things to be done:

        1) The last plate is changed into a variable/time dimension.
        Because a message mask is applied for plates only, the last
        axis of the mask must be applied to the message because the
        last plate is changed to a variable/time dimension.

        2) Because the message does not contain <X(n-1)X(n)> part,
        we'll put the last/third message to None meaning that it is
        empty.

        Parameters
        ----------
        index : int
            Index of the parent requesting the message.
        u_parents : list of list of ndarrays
            List of parents' moments.

        Returns
        -------
        m : list of ndarrays
            Message as a list of arrays.
        mask : boolean ndarray
            Mask telling which plates should be taken into account.
        """

        # Add the third empty message
        return [m_children[0], m_children[1], None]


# Make use of the converter
GaussianMarkovChainMoments.add_converter(GaussianMoments,
                                         _MarkovChainToGaussian)