<reponame>mariacer/cl_in_rnns
#!/usr/bin/env python3
# Copyright 2019 <NAME>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# title          :mnets/simple_rnn.py
# author         :mc, be
# contact        :mariacer, <EMAIL>
# created        :10/28/2019
# version        :1.0
# python_version :3.6.8
"""
SimpleRNN
---------

Implementation of a simple recurrent neural network that has stacked vanilla RNN
or LSTM layers that are optionally enclosed by fully-connected layers.

An example usage is as a main model, where the main weights are initialized
and protected by a method such as EWC, and the context-modulation patterns of
the neurons are produced by an external hypernetwork.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np

from mnets.mnet_interface import MainNetInterface
from utils.torch_utils import init_params

class SimpleRNN(nn.Module, MainNetInterface):
    """Implementation of a simple RNN.

    This is a simple recurrent network, that receives input vector
    :math:`\mathbf{x}` and outputs a vector :math:`\mathbf{y}` of real values.

    Note:
        The output is non-linear if the last layer is recurrent! Otherwise,
        logits are returned (cmp. attribute
        :attr:`mnets.mnet_interface.MainNetInterface.has_fc_out`).

    Attributes:
        bptt_depth (int): The truncation depth for backprop through time.
            If ``-1``, backprop through time (BPTT) will unroll all timesteps
            present in the input. Otherwise, the forward pass will detach the
            RNN hidden states smaller or equal than
            ``num_timesteps - bptt_depth`` timesteps, resulting in truncated
            BPTT (T-BPTT).
        num_rec_layers (int): Number of recurrent layers in this network (i.e.,
            length of constructor argument ``rnn_layers``).
        use_lstm (bool): See constructor argument ``use_lstm``.

    Args:
        n_in (int): Number of inputs.
        rnn_layers (list or tuple): List of integers. Each entry denotes the
            size of a recurrent layer. Recurrent layers will simply be stacked
            as layers of this network.

            If ``fc_layers_pre`` is empty, then the recurrent layers are the
            initial layers.
            If ``fc_layers`` is empty, then the last entry of this list will
            denote the output size.

            Note:
                This list may never be empty.
        fc_layers_pre (list or tuple): List of integers. Before the recurrent
            layers a set of fully-connected layers may be added. This might be
            specially useful when constructing recurrent autoencoders. The
            entries of this list will denote the sizes of those layers.

            If ``fc_layers_pre`` is not empty, its first entry will denote the
            input size of this network.
        fc_layers (list or tuple): List of integers. After the recurrent layers,
            a set of fully-connected layers is added. The entries of this list
            will denote the sizes of those layers.

            If ``fc_layers`` is not empty, its last entry will denote the output
            size of this network.
        activation: The nonlinearity used in hidden layers.
        use_lstm (bool): If set to `True``, the recurrent layers will be LSTM
            layers.
        use_bias (bool): Whether layers may have bias terms.
        no_weights (bool): If set to ``True``, no trainable parameters will be
            constructed, i.e., weights are assumed to be produced ad-hoc
            by a hypernetwork and passed to the :meth:`forward` method.
        init_weights (list, optional): This option is for convinience reasons.
            The option expects a list of parameter values that are used to
            initialize the network weights. As such, it provides a
            convinient way of initializing a network with a weight draw
            produced by the hypernetwork.

            Note, internal weights (see 
            :attr:`mnets.mnet_interface.MainNetInterface.weights`) will be
            affected by this argument only.
        kaiming_rnn_init (bool): By default, PyTorch initializes its recurrent
            layers uniformly with an interval defined by the square-root of the
            inverse of the layer size.

            If this option is enabled, then the recurrent layers will be
            initialized using the kaiming init as implemented by the function
            :func:`utils.torch_utils.init_params`.
        context_mod_last_step (bool): Whether context modulation is applied
            at the last time step os a recurrent layer only. If ``False``,
            context modulation is applied at every time step.

            Note:
                This option only applies if ``use_context_mod`` is ``True``.
        context_mod_num_ts (int, optional): The maximum number of timesteps.
            If specified, context-modulation with a different set of weights is
            applied at every timestep. If ``context_mod_separate_layers_per_ts``
            is ``True``, then a separate context-mod layer per timestep will be
            created. Otherwise, a single context-mod layer is created, but the
            expected parameter shapes for this layer are
            ``[context_mod_num_ts, *context_mod_shape]``.

            Note:
                This option only applies if ``use_context_mod`` is ``True``.
        context_mod_separate_layers_per_ts (bool): If specified, a separate
            context-mod layer per timestep is created (required if
            ``context_mod_no_weights`` is ``False``).
            
            Note:
                Only applies if ``context_mod_num_ts`` is specified.
        verbose (bool): Whether to print information (e.g., the number of
            weights) during the construction of the network.
        **kwargs: Keyword arguments regarding context modulation. This class
            can process the same context-modulation related arguments as class
            :class:`mnets.mlp.MLP` (plus the additional ones noted above).
    """
    def __init__(self, n_in=1, rnn_layers=(10,), fc_layers_pre=(),
                 fc_layers=(1,), activation=torch.nn.Tanh(), use_lstm=False,
                 use_bias=True, no_weights=False,
                 init_weights=None, kaiming_rnn_init=False,
                 context_mod_last_step=False,
                 context_mod_num_ts=-1,
                 context_mod_separate_layers_per_ts=False,
                 verbose=True,
                 **kwargs):
        # FIXME find a way using super to handle multiple inheritance.
        nn.Module.__init__(self)
        MainNetInterface.__init__(self)

        self._bptt_depth = -1

        # FIXME Arbitrary restriction.
        if activation is None or isinstance(activation, (torch.nn.ReLU, \
                torch.nn.Tanh)):
            self._a_fun = activation
        else:
            raise ValueError('Only linear, relu and tanh activations are ' + \
                             'allowed for recurrent networks.')

        if len(rnn_layers) == 0:
            raise ValueError('The network always needs to have at least one ' +
                             'recurrent layer.')
        if len(fc_layers) == 0:
            has_rec_out_layer = True
            #n_out = rnn_layers[-1]
        else:
            has_rec_out_layer = False
            #n_out = fc_layers[-1]

        self._n_in = n_in
        self._rnn_layers = rnn_layers
        self._fc_layers_pre = fc_layers_pre
        self._fc_layers = fc_layers

        self._no_weights = no_weights

        ### Parse or set context-mod arguments ###
        rem_kwargs = MainNetInterface._parse_context_mod_args(kwargs)
        if len(rem_kwargs) > 0:
            raise ValueError('Keyword arguments %s unknown.' % str(rem_kwargs))

        self._use_context_mod = kwargs['use_context_mod']
        self._context_mod_inputs = kwargs['context_mod_inputs']
        self._no_last_layer_context_mod = kwargs['no_last_layer_context_mod']
        self._context_mod_no_weights = kwargs['context_mod_no_weights']
        self._context_mod_post_activation = \
            kwargs['context_mod_post_activation']
        self._context_mod_gain_offset = kwargs['context_mod_gain_offset']
        self._context_mod_gain_softplus = kwargs['context_mod_gain_softplus']

        # Context-mod options specific to RNNs
        self._context_mod_last_step = context_mod_last_step
        # FIXME We have to specify this option even if
        # `context_mod_separate_layers_per_ts` is False (in order to set
        # sensible parameter shapes). However, the forward method can deal with
        # an arbitrary timestep length.
        self._context_mod_num_ts = context_mod_num_ts
        self._context_mod_separate_layers_per_ts = \
            context_mod_separate_layers_per_ts

        # More appropriate naming of option.
        self._context_mod_outputs = not self._no_last_layer_context_mod

        if context_mod_num_ts != -1:
            if context_mod_last_step:
                raise ValueError('Options "context_mod_last_step" and ' +
                                 '"context_mod_num_ts" are not compatible.')
            if not self._context_mod_no_weights and \
                    not context_mod_separate_layers_per_ts:
                raise ValueError('When applying context-mod per timestep ' +
                    'while maintaining weights internally, option' +
                    '"context_mod_separate_layers_per_ts" must be set.')
        ### Parse or set context-mod arguments - DONE ###

        self._has_bias = use_bias
        self._has_fc_out = True if not has_rec_out_layer else False
        # We need to make sure that the last 2 entries of `weights` correspond
        # to the weight matrix and bias vector of the last layer.
        self._mask_fc_out = True if not has_rec_out_layer else False
        # Note, recurrent layers always use non-linearities and their activities
        # are squashed by a non-linearity (otherwise, internal states could
        # vanish/explode with increasing sequence length).
        self._has_linear_out = True if not has_rec_out_layer else False

        self._param_shapes = []
        self._param_shapes_meta = []

        self._weights = None if no_weights and self._context_mod_no_weights \
            else nn.ParameterList()
        self._hyper_shapes_learned = None \
            if not no_weights and not self._context_mod_no_weights else []
        self._hyper_shapes_learned_ref = None if self._hyper_shapes_learned \
            is None else []

        self._layer_weight_tensors = nn.ParameterList()
        self._layer_bias_vectors = nn.ParameterList()

        self._use_lstm = use_lstm
        if use_lstm:
            self._rnn_fct = self.lstm_rnn_step
        else:
            self._rnn_fct = self.basic_rnn_step

        #################################################
        ### Define and initialize context mod weights ###
        #################################################

        # The context-mod layers consist of sequential layers ordered as:
        # - initial fully-connected layers if len(fc_layers_pre)>0
        # - recurrent layers
        # - final fully-connected layers if len(fc_layers)>0

        self._context_mod_layers = nn.ModuleList() if self._use_context_mod \
            else None

        self._cm_rnn_start_ind = 0
        self._num_fc_cm_layers = None
        if self._use_context_mod:
            cm_layer_inds = []
            cm_shapes = []

            # Gather sizes of all activation vectors within the network that
            # will be subject to context-modulation.
            if self._context_mod_inputs:
                self._cm_rnn_start_ind += 1
                cm_shapes.append([n_in])

                # We reserve layer zero for input context-mod. Otherwise, there
                # is no layer zero.
                cm_layer_inds.append(0)

            if len(fc_layers_pre) > 0:
                self._cm_rnn_start_ind += len(fc_layers_pre)

            # We use odd numbers for actual layers and even number for all
            # context-mod layers.
            rem_cm_inds = range(2, 2*(len(fc_layers_pre)+len(rnn_layers)+\
                len(fc_layers))+1, 2)

            num_rec_cm_layers = len(rnn_layers)
            if has_rec_out_layer and not self._context_mod_outputs:
                num_rec_cm_layers -= 1
            self._num_rec_cm_layers = num_rec_cm_layers

            jj = 0
            # Add initial fully-connected context-mod layers.
            num_fc_pre_cm_layers = len(fc_layers_pre)
            self._num_fc_pre_cm_layers = num_fc_pre_cm_layers
            for i in range(num_fc_pre_cm_layers):
                cm_shapes.append([fc_layers_pre[i]])
                cm_layer_inds.append(rem_cm_inds[jj])
                jj += 1

            # Add recurrent context-mod layers.
            for i in range(num_rec_cm_layers):
                if context_mod_num_ts != -1:
                    if context_mod_separate_layers_per_ts:
                        cm_rnn_shapes = [[rnn_layers[i]]] * context_mod_num_ts
                    else:
                        # Only a single context-mod layer will be added, but we
                        # directly edit the correponding `param_shape` later.
                        assert self._context_mod_no_weights
                        cm_rnn_shapes = [[rnn_layers[i]]]
                else:
                    cm_rnn_shapes = [[rnn_layers[i]]]

                cm_shapes.extend(cm_rnn_shapes)
                cm_layer_inds.extend([rem_cm_inds[jj]] * len(cm_rnn_shapes))
                jj += 1

            # Add final fully-connected context-mod layers.
            num_fc_cm_layers = len(fc_layers)
            if num_fc_cm_layers > 0 and not self._context_mod_outputs:
                num_fc_cm_layers -= 1
            self._num_fc_cm_layers = num_fc_cm_layers
            for i in range(num_fc_cm_layers):
                cm_shapes.append([fc_layers[i]])
                cm_layer_inds.append(rem_cm_inds[jj])
                jj += 1

            self._add_context_mod_layers(cm_shapes, cm_layers=cm_layer_inds)

            if context_mod_num_ts != -1 and not \
                    context_mod_separate_layers_per_ts:
                # In this case, there is only one context-mod layer for each
                # recurrent layer, but we want to have separate weights per
                # timestep.
                # Hence, we adapt the expected parameter shape, such that we
                # get a different set of weights per timestep. This will be
                # split into multiple weights that are succesively fed into the
                # same layer inside the forward method.

                for i in range(num_rec_cm_layers):
                    cmod_layer = \
                        self.context_mod_layers[self._cm_rnn_start_ind+i]
                    cm_shapes_rnn = [[context_mod_num_ts, *s] for s in \
                                      cmod_layer.param_shapes]

                    ps_ind = int(np.sum([ \
                        len(self.context_mod_layers[ii].param_shapes) \
                        for ii in range(self._cm_rnn_start_ind+i)]))
                    self._param_shapes[ps_ind:ps_ind+len(cm_shapes_rnn)] = \
                        cm_shapes_rnn
                    assert self._hyper_shapes_learned is not None
                    self._hyper_shapes_learned[ \
                        ps_ind:ps_ind+len(cm_shapes_rnn)] = cm_shapes_rnn

        ########################
        ### Internal weights ###
        ########################
        prev_dim = self._n_in

        def define_fc_layer_weights(fc_layers, prev_dim, num_prev_layers):
            """Define the weights and shapes of the fully-connected layers.

            Args:
                fc_layers (list): The list of fully-connected layer dimensions.
                prev_dim (int): The output size of the previous layer.
                num_prev_layers (int): The number of upstream layers to the 
                    current one (a layer with its corresponding
                    context-mod layer(s) count as one layer). Count should
                    start at ``1``.

            Returns:
                (int): The output size of the last fully-connected layer
                considered here.
            """
            # FIXME We should instead build an MLP instance. But then we still
            # have to adapt all attributes accordingly.
            for i, n_fc in enumerate(fc_layers):
                s_w = [n_fc, prev_dim]
                s_b = [n_fc] if self._has_bias else None

                for j, s in enumerate([s_w, s_b]):
                    if s is None:
                        continue

                    is_bias = True
                    if j % 2 == 0:
                        is_bias = False

                    if not self._no_weights:
                        self._weights.append(nn.Parameter(torch.Tensor(*s),
                                                          requires_grad=True))
                        if is_bias:
                            self._layer_bias_vectors.append(self._weights[-1])
                        else:
                            self._layer_weight_tensors.append(self._weights[-1])
                    else:
                        self._hyper_shapes_learned.append(s)
                        self._hyper_shapes_learned_ref.append( \
                            len(self.param_shapes))

                    self._param_shapes.append(s)
                    self._param_shapes_meta.append({
                        'name': 'bias' if is_bias else 'weight',
                        'index': -1 if self._no_weights else \
                            len(self._weights)-1,
                        'layer': i * 2 + num_prev_layers, # Odd numbers
                    })

                prev_dim = n_fc

            return prev_dim

        ### Initial fully-connected layers.
        prev_dim = define_fc_layer_weights(self._fc_layers_pre, prev_dim, 1)

        ### Recurrent layers.
        coeff = 4 if self._use_lstm else 1
        for i, n_rec in enumerate(self._rnn_layers):
            # Input-to-hidden
            s_w_ih = [n_rec*coeff, prev_dim]
            s_b_ih = [n_rec*coeff] if use_bias else None

            # Hidden-to-hidden
            s_w_hh = [n_rec*coeff, n_rec]
            s_b_hh = [n_rec*coeff] if use_bias else None

            # Hidden-to-output.
            # Note, for an LSTM cell, the hidden state vector is also the
            # output vector.
            if not self._use_lstm:
                s_w_ho = [n_rec, n_rec]
                s_b_ho = [n_rec] if use_bias else None
            else:
                s_w_ho = None
                s_b_ho = None

            for j, s in enumerate([s_w_ih, s_b_ih, s_w_hh, s_b_hh, s_w_ho,
                                   s_b_ho]):
                if s is None:
                    continue

                is_bias = True
                if j % 2 == 0:
                    is_bias = False

                wtype = 'ih'
                if 2 <= j < 4:
                    wtype = 'hh'
                elif j >=4:
                    wtype = 'ho'

                if not no_weights:
                    self._weights.append(nn.Parameter(torch.Tensor(*s),
                                                      requires_grad=True))
                    if is_bias:
                        self._layer_bias_vectors.append(self._weights[-1])
                    else:
                        self._layer_weight_tensors.append(self._weights[-1])
                else:
                    self._hyper_shapes_learned.append(s)
                    self._hyper_shapes_learned_ref.append( \
                        len(self.param_shapes))

                self._param_shapes.append(s)
                self._param_shapes_meta.append({
                    'name': 'bias' if is_bias else 'weight',
                    'index': -1 if no_weights else len(self._weights)-1,
                    'layer': i * 2 + 1 + 2 * len(fc_layers_pre), # Odd numbers
                    'info': wtype
                })

            prev_dim = n_rec

        ### Fully-connected layers.
        prev_dim = define_fc_layer_weights(self._fc_layers, prev_dim, \
            1 + 2 * len(fc_layers_pre) + 2 * len(rnn_layers))

        ### Initialize weights.
        if init_weights is not None:
            assert self._weights is not None
            assert len(init_weights) == len(self.weights)
            for i in range(len(init_weights)):
                assert np.all(np.equal(list(init_weights[i].shape),
                                       self.weights[i].shape))
                self.weights[i].data = init_weights[i]
        else:
            rec_start = len(fc_layers_pre)
            rec_end = rec_start + len(rnn_layers) * (2 if use_lstm else 3)
            # Note, Pytorch applies a uniform init to its recurrent layers, as
            # defined here:
            # https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py#L155
            for i in range(len(self._layer_weight_tensors)):
                if i >=rec_start and i < rec_end:
                    # Recurrent layer weights.
                    if kaiming_rnn_init:
                        init_params(self._layer_weight_tensors[i],
                            self._layer_bias_vectors[i] if use_bias else None)
                    else:
                        a = 1.0 / math.sqrt(rnn_layers[(i-rec_start) // \
                            (2 if use_lstm else 3)])
                        nn.init.uniform_(self._layer_weight_tensors[i], -a, a)
                        if use_bias:
                            nn.init.uniform_(self._layer_bias_vectors[i], -a, a)
                else:
                    # FC layer weights.
                    init_params(self._layer_weight_tensors[i],
                        self._layer_bias_vectors[i] if use_bias else None)

        num_weights = MainNetInterface.shapes_to_num_weights(self._param_shapes)

        if verbose:
            if self._use_context_mod:
                cm_num_weights =  \
                    MainNetInterface.shapes_to_num_weights(cm_shapes)

            print('Creating a simple RNN with %d weights' % num_weights
                  + (' (including %d weights associated with-' % cm_num_weights
                     + 'context modulation)' if self._use_context_mod else '')
                  + '.')

        self._is_properly_setup()

    @property
    def bptt_depth(self):
        """Getter for attribute :attr:`bptt_depth`."""
        return self._bptt_depth

    @bptt_depth.setter
    def bptt_depth(self, value):
        """Setter for attribute :attr:`bptt_depth`."""
        self._bptt_depth = value

    @property
    def num_rec_layers(self):
        """Getter for read-only attribute :attr:`num_rec_layers`."""
        return len(self._rnn_layers)

    @property
    def use_lstm(self):
        """Getter for read-only attribute :attr:`use_lstm`."""
        return self._use_lstm

    def split_cm_weights(self, cm_weights, condition, num_ts=0):
        """Split context-mod weights per context-mod layer.

        Args:
            cm_weights (torch.Tensor): All context modulation weights.
            condition (optional, int): If provided, then this argument will be
                passed as argument ``ckpt_id`` to the method
                :meth:`utils.context_mod_layer.ContextModLayer.forward`.
            num_ts (int): The length of the sequences.

        Returns:
            (Tuple): Where the tuple contains:

            - **cm_inputs_weights**: The cm input weights.
            - **cm_fc_pre_layer_weights**: The cm pre-recurrent weights.
            - **cm_rec_layer_weights**: The cm recurrent weights.
            - **cm_fc_layer_weights**: The cm post-recurrent weights.
            - **n_cm_rec**: The number of recurrent cm layers.
            - **cmod_cond**: The context-mod condition.
        """

        n_cm_rec = -1
        cm_fc_pre_layer_weights = None
        cm_fc_layer_weights = None
        cm_inputs_weights = None
        cm_rec_layer_weights = None
        if cm_weights is not None:
            if self._context_mod_num_ts != -1 and \
                    self._context_mod_separate_layers_per_ts:
                assert num_ts <= self._context_mod_num_ts

            # Note, an mnet layer might contain multiple context-mod layers
            # (a recurrent layer can have a separate context-mod layer per
            # timestep).
            cm_fc_pre_layer_weights = []
            cm_rec_layer_weights = [[] for _ in range(self._num_rec_cm_layers)]
            cm_fc_layer_weights = []

            # Number of cm-layers per recurrent layer.
            n_cm_per_rec = self._context_mod_num_ts if \
                self._context_mod_num_ts != -1 and \
                    self._context_mod_separate_layers_per_ts else 1
            n_cm_rec = n_cm_per_rec * self._num_rec_cm_layers

            cm_start = 0
            for i, cm_layer in enumerate(self.context_mod_layers):
                cm_end = cm_start + len(cm_layer.param_shapes)

                if i == 0 and self._context_mod_inputs:
                    cm_inputs_weights = cm_weights[cm_start:cm_end]
                elif i < self._cm_rnn_start_ind:
                    cm_fc_pre_layer_weights.append(cm_weights[cm_start:cm_end])
                elif i >= self._cm_rnn_start_ind and \
                        i < self._cm_rnn_start_ind + n_cm_rec:
                    # Index of recurrent layer.
                    i_r = (i-self._cm_rnn_start_ind) // n_cm_per_rec
                    cm_rec_layer_weights[i_r].append( \
                        cm_weights[cm_start:cm_end])
                else:
                    cm_fc_layer_weights.append(cm_weights[cm_start:cm_end])
                cm_start = cm_end

            # We need to split the context-mod weights in the following case,
            # as they are currently just stacked on top of each other.
            if self._context_mod_num_ts != -1 and \
                    not self._context_mod_separate_layers_per_ts:
                for i, cm_w_list in enumerate(cm_rec_layer_weights):
                    assert len(cm_w_list) == 1

                    cm_rnn_weights = cm_w_list[0]
                    cm_rnn_layer = self.context_mod_layers[ \
                        self._cm_rnn_start_ind+i]

                    assert len(cm_rnn_weights) == len(cm_rnn_layer.param_shapes)
                    # The first dimension are the weights of this layer per
                    # timestep.
                    num_ts_cm = -1
                    for j, s in enumerate(cm_rnn_layer.param_shapes):
                        assert len(cm_rnn_weights[j].shape) == len(s) + 1
                        if j == 0:
                            num_ts_cm = cm_rnn_weights[j].shape[0]
                        else:
                            assert num_ts_cm == cm_rnn_weights[j].shape[0]
                    assert num_ts <= num_ts_cm

                    cm_w_chunked = [None] * len(cm_rnn_weights)
                    for j, cm_w in enumerate(cm_rnn_weights):
                        cm_w_chunked[j] = torch.chunk(cm_w, num_ts_cm, dim=0)

                    # Now we gather all these chunks to assemble the weights
                    # needed per timestep (as if
                    # `_context_mod_separate_layers_per_t` were True).
                    cm_w_list = []
                    for j in range(num_ts_cm):
                        tmp_list = []
                        for chunk in cm_w_chunked:
                            tmp_list.append(chunk[j].squeeze(dim=0))
                        cm_w_list.append(tmp_list)
                    cm_rec_layer_weights[i] = cm_w_list

            # Note, the last layer does not necessarily have context-mod
            # (depending on `self._context_mod_outputs`).
            if len(cm_rec_layer_weights) < len(self._rnn_layers):
                cm_rec_layer_weights.append(None)
            if len(cm_fc_layer_weights) < len(self._fc_layers):
                cm_fc_layer_weights.append(None)


        #######################
        ### Parse condition ###
        #######################
        cmod_cond = None
        if condition is not None:
            assert isinstance(condition, int)
            cmod_cond = condition

            # Note, the cm layer will ignore the cmod condition if weights
            # are passed.
            # FIXME Find a more elegant solution.
            cm_inputs_weights = None
            cm_fc_pre_layer_weights = [None] * len(cm_fc_pre_layer_weights)
            cm_rec_layer_weights = [[None] * len(cm_ws) for cm_ws in \
                                    cm_rec_layer_weights]
            cm_fc_layer_weights = [None] * len(cm_fc_layer_weights)

        return cm_inputs_weights, cm_fc_pre_layer_weights, cm_fc_layer_weights,\
            cm_rec_layer_weights, n_cm_rec, cmod_cond

    def split_internal_weights(self, int_weights):
        """Split internal weights per layer.

        Args:
            int_weights (torch.Tensor): All internal weights.

        Returns:
            (Tuple): Where the tuple contains:

            - **fc_pre_w_weights**: The pre-recurrent w weights.
            - **fc_pre_b_weights**: The pre-recurrent b weights.
            - **rec_weights**: The recurrent weights.
            - **fc_w_weights**:The post-recurrent w weights.
            - **fc_b_weights**: The post-recurrent b weights.
        """
        n_cm = self._num_context_mod_shapes()

        int_meta = self.param_shapes_meta[n_cm:]
        assert len(int_meta) == len(int_weights)
        fc_pre_w_weights = []
        fc_pre_b_weights = []
        rec_weights =[[] for _ in range(len(self._rnn_layers))]
        fc_w_weights = []
        fc_b_weights = []

        # Number of pre-fc weights in total.
        n_fc_pre = len(self._fc_layers_pre)
        if self.has_bias:
            n_fc_pre *= 2

        # Number of weights per recurrent layer.
        if self._use_lstm:
            n_rw = 4 if self.has_bias else 2
        else:
            n_rw = 6 if self.has_bias else 3

        for i, w in enumerate(int_weights):
            if i < n_fc_pre: # fc pre weights
                if int_meta[i]['name'] == 'weight':
                    fc_pre_w_weights.append(w)
                else:
                    assert int_meta[i]['name'] == 'bias'
                    fc_pre_b_weights.append(w)
            elif i >= n_fc_pre and \
                    i < n_rw * len(self._rnn_layers) + n_fc_pre: # recurrent w
                r_ind = (i - n_fc_pre) // n_rw
                rec_weights[r_ind].append(w)
            else: # fc weights
                if int_meta[i]['name'] == 'weight':
                    fc_w_weights.append(w)
                else:
                    assert int_meta[i]['name'] == 'bias'
                    fc_b_weights.append(w)

        if not self.has_bias:
            assert len(fc_pre_b_weights) == 0
            fc_pre_b_weights = [None] * len(fc_pre_w_weights)

            assert len(fc_b_weights) == 0
            fc_b_weights = [None] * len(fc_w_weights)

        return fc_pre_w_weights, fc_pre_b_weights, rec_weights, fc_w_weights, \
            fc_b_weights

    def split_weights(self, weights):
        """Split weights into internal and context-mod weights.

        Extract which weights should be used,  I.e., are we using internally
        maintained weights or externally given ones or are we even mixing
        between these groups.

        Args:
            weights (torch.Tensor): All weights.

        Returns:
            (Tuple): Where the tuple contains:

            - **int_weights**: The internal weights.
            - **cm_weights**: The context-mod weights.
        """
        n_cm = self._num_context_mod_shapes()

        ### FIXME Code copied from MLP its `forward` method ###

        # Make sure cm_weights are either `None` or have the correct dimensions.
        if weights is None:
            weights = self.weights

            if self._use_context_mod:
                cm_weights = weights[:n_cm]
                int_weights = weights[n_cm:]
            else:
                cm_weights = None
                int_weights = weights
        else:
            int_weights = None
            cm_weights = None

            if isinstance(weights, dict):
                assert 'internal_weights' in weights.keys() or \
                       'mod_weights' in weights.keys()
                if 'internal_weights' in weights.keys():
                    int_weights = weights['internal_weights']
                if 'mod_weights' in weights.keys():
                    cm_weights = weights['mod_weights']
            else:
                if self._use_context_mod and \
                        len(weights) == n_cm:
                    cm_weights = weights
                else:
                    assert len(weights) == len(self.param_shapes)
                    if self._use_context_mod:
                        cm_weights = weights[:n_cm]
                        int_weights = weights[n_cm:]
                    else:
                        int_weights = weights

            if self._use_context_mod and cm_weights is None:
                if self._context_mod_no_weights:
                    raise Exception('Network was generated without weights ' +
                        'for context-mod layers. Hence, they must be passed ' +
                        'via the "weights" option.')
                cm_weights = self.weights[:n_cm]
            if int_weights is None:
                if self._no_weights:
                    raise Exception('Network was generated without internal ' +
                        'weights. Hence, they must be passed via the ' +
                        '"weights" option.')
                if self._context_mod_no_weights:
                    int_weights = self.weights
                else:
                    int_weights = self.weights[n_cm:]

            # Note, context-mod weights might have different shapes, as they
            # may be parametrized on a per-sample basis.
            if self._use_context_mod:
                assert len(cm_weights) == n_cm
            int_shapes = self.param_shapes[n_cm:]
            assert len(int_weights) == len(int_shapes)
            for i, s in enumerate(int_shapes):
                assert np.all(np.equal(s, list(int_weights[i].shape)))

        ### FIXME Code copied until here ###
        return int_weights, cm_weights

    def forward(self, x, weights=None, distilled_params=None, condition=None,
                return_hidden=False, return_hidden_int=False):
        """Compute the output :math:`y` of this network given the input
        :math:`x`.

        Args:
            (....): See docstring of method
                :meth:`mnets.mnet_interface.MainNetInterface.forward`. We
                provide some more specific information below.
            weights (list or dict): See argument ``weights`` of method
                :meth:`mnets.mlp.MLP.forward`.
            condition (optional, int): If provided, then this argument will be
                passed as argument ``ckpt_id`` to the method
                :meth:`utils.context_mod_layer.ContextModLayer.forward`.
            return_hidden (bool, optional): If ``True``, all hidden activations
                of fully-connected and recurrent layers (where we defined
                :math:`y_t` as hidden state of vannila RNN layers as these are
                the layer outputs passed to the next layer) are returned.
                recurrent hidden activations will be returned.
            return_hidden_int (bool, optional): If ``True``, in addition to
                ``hidden``, an additional variable ``hidden_int`` is returned
                containing the internal hidden states of recurrent layers (i.e.,
                the cell states :math:`c_t` for LSTMs and the actual hidden
                state :math:`h_t` for Elman layers) are returned.

        Returns:
            (torch.Tensor or tuple): Where the tuple is containing:

            - **output** (torch.Tensor): The output of the network.
            - **hidden** (list): If ``return_hidden`` is ``True``, then the
              hidden activities of each layer are returned, which have the shape
              ``(seq_length, batch_size, n_hidden)``.
            - **hidden_int**: If ``return_hidden_int`` is ``True``, then in
              addition to ``hidden`` a tensor ``hidden_int`` is returned
              containing internal hidden states of recurrent layers.
        """
        assert distilled_params is None

        if ((not self._use_context_mod and self._no_weights) or \
                (self._no_weights or self._context_mod_no_weights)) and \
                weights is None:
            raise Exception('Network was generated without weights. ' +
                            'Hence, "weights" option may not be None.')

        #######################
        ### Extract weights ###
        #######################
        # Extract which weights should be used.
        int_weights, cm_weights = self.split_weights(weights)

        ### Split context-mod weights per context-mod layer.
        cm_inputs_weights, cm_fc_pre_layer_weights, cm_fc_layer_weights, \
            cm_rec_layer_weights, n_cm_rec, cmod_cond = self.split_cm_weights(
                cm_weights, condition, num_ts=x.shape[0])

        ### Extract internal weights.
        fc_pre_w_weights, fc_pre_b_weights, rec_weights, fc_w_weights, \
            fc_b_weights = self.split_internal_weights(int_weights)

        ###########################
        ### Forward Computation ###
        ###########################
        ret_hidden = None
        if return_hidden:
            ret_hidden = []

        h = x

        cm_offset = 0
        if self._use_context_mod and self._context_mod_inputs:
            cm_offset += 1
            # Apply context modulation in the inputs.
            h = self._context_mod_layers[0].forward(h,
                weights=cm_inputs_weights, ckpt_id=cmod_cond, bs_dim=1)

        ### Initial fully-connected layer activities.
        ret_hidden, h = self.compute_fc_outputs(h, fc_pre_w_weights, \
            fc_pre_b_weights, len(self._fc_layers_pre), \
            cm_fc_pre_layer_weights, cm_offset, cmod_cond, False, ret_hidden)

        ### Recurrent layer activities.
        ret_hidden_int = [] # the internal hidden activations
        for d in range(len(self._rnn_layers)):
            if self._use_context_mod:
                h, h_int = self.compute_hidden_states(h, d, rec_weights[d],
                    cm_rec_layer_weights[d], cmod_cond)
            else:
                h, h_int = self.compute_hidden_states(h, d, rec_weights[d],
                                                      None, None)
            if ret_hidden is not None:
                ret_hidden.append(h)
                ret_hidden_int.append(h_int)

        ### Fully-connected layer activities.
        cm_offset = self._cm_rnn_start_ind + n_cm_rec
        ret_hidden, h = self.compute_fc_outputs(h, fc_w_weights, fc_b_weights, \
            self._num_fc_cm_layers, cm_fc_layer_weights, cm_offset, cmod_cond,
            True, ret_hidden)

        # FIXME quite ugly
        if return_hidden:
            # The last element is the output activity.
            ret_hidden.pop()
            if return_hidden_int:
                return h, ret_hidden, ret_hidden_int
            else:
                return h, ret_hidden
        else:
            return h

    def compute_fc_outputs(self, h, fc_w_weights, fc_b_weights, num_fc_cm_layers,
                cm_fc_layer_weights, cm_offset, cmod_cond, is_post_fc, 
                ret_hidden):
        """Compute the forward pass through the fully-connected layers.

        This method also appends activations to ``ret_hidden``.

        Args:
            h (torch.Tensor): The input from the previous layer.
            fc_w_weights (list): The weights for the fc layers.
            fc_b_weights (list): The biases for the fc layers.
            num_fc_cm_layers (int): The number of context-modulation
                layers associated with this set of fully-connected layers.
            cm_fc_layer_weights (list): The context-modulation weights
                associated with the current layers.
            cm_offset (int): The index to access the correct context-mod
                layers.
            cmod_cond (bool): Some condition to perform context modulation.
            is_post_fc (bool); Whether those layers are applied as last
                layers of the network. In this case, there will be no
                activation applied to the last layer outputs.
            ret_hidden (list or None): The hidden recurrent activations.

        Return:
            (Tuple): Tuple containing:

            - **ret_hidden**: The hidden recurrent activations.
            - **h**: Transformed activation ``h``.
        """
        for d in range(len(fc_w_weights)):
            use_cm = self._use_context_mod and d < num_fc_cm_layers
            # Compute output.
            h = F.linear(h, fc_w_weights[d], bias=fc_b_weights[d])

            # Context-dependent modulation (pre-activation).
            if use_cm and not self._context_mod_post_activation:
                h = self._context_mod_layers[cm_offset+d].forward(h,
                    weights=cm_fc_layer_weights[d], ckpt_id=cmod_cond, 
                    bs_dim=1)

            # Non-linearity
            # Note, non-linearity is not applied to outputs of the network.
            if self._a_fun is not None and \
                    (not is_post_fc or d < len(fc_w_weights)-1):
                h = self._a_fun(h)

            # Context-dependent modulation (post-activation).
            if use_cm and self._context_mod_post_activation:
                h = self._context_mod_layers[cm_offset+d].forward(h,
                    weights=cm_fc_layer_weights[d], ckpt_id=cmod_cond, 
                    bs_dim=1)

            if ret_hidden is not None:
                ret_hidden.append(h)

        return ret_hidden, h

    def compute_hidden_states(self, x, layer_ind, int_weights, cm_weights,
                              ckpt_id, h_0=None, c_0=None):
        """Compute the hidden states for the recurrent layer ``layer_ind`` from
        a sequence of inputs :math:`x`.

        If so specified, context modulation is applied before or after the
        nonlinearities.

        Args:
            x: The inputs :math:`x` to the layer. :math:`x` has shape
                ``[sequence_len, batch_size, n_hidden_prev]``.
            layer_ind (int): Index of the layer.
            int_weights: Internal weights associated with this recurrent layer.
            cm_weights: Context modulation weights.
            ckpt_id: Will be passed as option ``ckpt_id`` to method
                :meth:`utils.context_mod_layer.ContextModLayer.forward` if
                context-mod layers are used.
            h_0 (torch.Tensor, optional): The initial state for :math:`h`.
            c_0 (torch.Tensor, optional): The initial state for :math:`c`. Note
                that for LSTMs, if the initial state is to be defined, this
                variable is necessary also, not only :math:`h_0`, whereas for
                vanilla RNNs it is enough to provide :math:`h_0` as :math:`c_0`
                represents the output of the layer and it can be easily computed
                from `h_0`.

        Returns:
            (tuple): Tuple containing:

            - **outputs** (torch.Tensor): The sequence of visible hidden states
              given the input. It has shape 
              ``[sequence_len, batch_size, n_hidden]``.
            - **hiddens** (torch.Tensor): The sequence of hidden states given
              the input. For LSTMs, this corresponds to :math:`c`.
              It has shape ``[sequence_len, batch_size, n_hidden]``.
        """
        seq_length, batch_size, n_hidden_prev = x.shape
        n_hidden = self._rnn_layers[layer_ind]

        # Generate initial hidden states.
        # Note that for LSTMs h_0 is the hidden state and output vector whereas
        # c_0 is the internal cell state vector.
        # For a vanilla RNN h_0 is the hidden state whereas c_0 is the output
        # vector.
        if h_0 is None:
            h_0 = (torch.zeros(batch_size, n_hidden, device=x.device))
        if c_0 is None:
            c_0 = (torch.zeros(batch_size, n_hidden, device=x.device))
        assert h_0.shape[0] == c_0.shape[0] == batch_size
        assert h_0.shape[1] == c_0.shape[1] == n_hidden

        # If we want to apply context modulation in each time step, we need
        # to split the input sequence and call pytorch function at every
        # time step.
        outputs = []
        hiddens = []
        h_t = h_0
        c_t = c_0
        for t in range(seq_length):
            x_t = x[t,:,:]

            if cm_weights is not None and self._context_mod_num_ts != -1:
                curr_cm_weights = cm_weights[t]
            elif cm_weights is not None:
                assert len(cm_weights) == 1
                curr_cm_weights = cm_weights[0]
            else:
                curr_cm_weights = cm_weights

            # Compute the actual rnn step (either vanilla or LSTM, depending on
            # the flag self._use_lstm).
            is_last_step = t==(seq_length-1)
            h_t, c_t = self._rnn_fct(layer_ind, t, x_t, (h_t, c_t), int_weights,
                                     curr_cm_weights, ckpt_id, is_last_step)

            if self.bptt_depth != -1:
                if t < (seq_length - self.bptt_depth):
                    # Detach hidden/output states, such that we don't backprop
                    # through these timesteps.
                    h_t = h_t.detach()
                    c_t = c_t.detach()

            # FIXME Solution is a bit ugly. For an LSTM, the hidden state is
            # also the output whereas a normal RNN has a separate output.
            if self._use_lstm:
                outputs.append(h_t)
                hiddens.append(c_t)
            else:
                outputs.append(c_t)
                hiddens.append(h_t)

        return torch.stack(outputs), torch.stack(hiddens)


    def basic_rnn_step(self, d, t, x_t, h_t, int_weights, cm_weights, ckpt_id,
                       is_last_step):
        """Perform vanilla rnn pass from inputs to hidden units.

        Apply context modulation if necessary (i.e. if ``cm_weights`` is
        not ``None``).

        This function implements a step of an
        `Elman RNN <https://en.wikipedia.org/wiki/\
Recurrent_neural_network#Elman_networks_and_Jordan_networks>`__.

        Note:
            We made the following design choice regarding context-modulation.
            In contrast to the LSTM, the Elman network layer consists of "two
            steps", updating the hidden state and computing an output based
            on this hidden state. To be fair, context-mod should influence both
            these "layers". Therefore, we apply context-mod twice, but using the
            same weights. This of course assumes that the hidden state and
            output vector have the same dimensionality.

        Args:
            d (int): Index of the layer.
            t (int): Current timestep.
            x_t: Tensor of size ``[batch_size, n_hidden_prev]`` with inputs.
            h_t (tuple): Tuple of length 2, containing two tensors of size
                ``[batch_size, n_hidden]`` with previous hidden states ``h`` and
                and previous outputs ``y``.

                Note:
                    The previous outputs ``y`` are ignored by this method, since
                    they are not required in an Elman RNN step.
            int_weights: See docstring of method :meth:`compute_hidden_states`.
            cm_weights (list): The weights of the context-mod layer, if context-
                mod should be applied.
            ckpt_id: See docstring of method :meth:`compute_hidden_states`.
            is_last_step (bool): Whether the current time step is the last one.

        Returns:
            (tuple): Tuple containing:

            - **h_t** (torch.Tensor): The tensor ``h_t`` of size
              ``[batch_size, n_hidden]`` with the new hidden state.
            - **y_t** (torch.Tensor): The tensor ``y_t`` of size
              ``[batch_size, n_hidden]`` with the new cell state.
        """
        h_t = h_t[0]

        use_cm = self._use_context_mod and d < self._num_rec_cm_layers
        # Determine the index of the hidden context mod layer.
        # Number of cm-layers per recurrent layer.
        n_cm_per_rec = self._context_mod_num_ts if \
            self._context_mod_num_ts != -1 and \
                self._context_mod_separate_layers_per_ts else 1
        cm_idx = self._cm_rnn_start_ind + d * n_cm_per_rec
        if self._context_mod_num_ts != -1 and \
                self._context_mod_separate_layers_per_ts:
            cm_idx += t
 
        if self.has_bias:
            assert len(int_weights) == 6
            weight_ih = int_weights[0]
            bias_ih = int_weights[1]
            weight_hh = int_weights[2]
            bias_hh = int_weights[3]
        else:
            assert len(int_weights) == 3
            weight_ih = int_weights[0]
            bias_ih = None
            weight_hh = int_weights[1]
            bias_hh = None

        ###########################
        ### Update hidden state ###
        ###########################
        h_t = x_t @ weight_ih.t() + h_t @ weight_hh.t()
        if self.has_bias:
            h_t += bias_ih + bias_hh

        # Context-dependent modulation (pre-activation).
        if use_cm and not self._context_mod_post_activation:
            # Only apply context mod if you are in the last time step, or if
            # you want to apply it in every single time step (i.e. if
            # self._context_mod_last_step is False).
            if not self._context_mod_last_step or is_last_step:
                h_t = self._context_mod_layers[cm_idx].forward(h_t,
                    weights=cm_weights, ckpt_id=ckpt_id)

        # Compute activation.
        if self._a_fun is not None:
            h_t = self._a_fun(h_t)

        # Context-dependent modulation (post-activation).
        if use_cm and self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                h_t = self._context_mod_layers[cm_idx].forward(h_t,
                    weights=cm_weights, ckpt_id=ckpt_id)

        ######################
        ### Compute output ###
        ######################
        y_t = self.compute_basic_rnn_output(h_t, int_weights, use_cm, 
            cm_weights, cm_idx, ckpt_id, is_last_step)

        return h_t, y_t

    def compute_basic_rnn_output(self, h_t, int_weights, use_cm, cm_weights,
                                 cm_idx, ckpt_id, is_last_step):
        """Compute the output of a vanilla RNN given the hidden state.

        Args:
            (...): See docstring of method :meth:`basic_rnn_step`.
            use_cm (boolean): Whether context modulation is being used.
            cm_idx (int): Index of the context-mod layer.

        Returns:
            (torch.tensor): The output.
        """
        if self.has_bias:
            weight_ho = int_weights[4]
            bias_ho = int_weights[5]
        else:
            weight_ho = int_weights[2]
            bias_ho = None

        y_t = h_t @ weight_ho.t()
        if self.has_bias:
            y_t += bias_ho

        # Context-dependent modulation (pre-activation).
        if use_cm and not self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                y_t = self._context_mod_layers[cm_idx].forward(y_t,
                    weights=cm_weights, ckpt_id=ckpt_id)

        # Compute activation.
        if self._a_fun is not None:
            y_t = self._a_fun(y_t)

        # Context-dependent modulation (post-activation).
        if use_cm and self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                y_t = self._context_mod_layers[cm_idx].forward(y_t,
                    weights=cm_weights, ckpt_id=ckpt_id)
        return y_t

    def lstm_rnn_step(self, d, t, x_t, h_t, int_weights, cm_weights, ckpt_id,
                      is_last_step):
        """ Perform an LSTM pass from inputs to hidden units.

        Apply masks to the temporal sequence for computing the loss.
        Obtained from:

            https://mlexplained.com/2019/02/15/building-an-lstm-from-scratch-\
in-pytorch-lstms-in-depth-part-1/

        and:

            https://d2l.ai/chapter_recurrent-neural-networks/lstm.html

        Args:
            d (int): Index of the layer.
            t (int): Current timestep.
            x_t: Tensor of size ``[batch_size, n_inputs]`` with inputs.
            h_t (tuple): Tuple of length 2, containing two tensors of size
                ``[batch_size, n_hidden]`` with previous hidden states ``h`` and
                ``c``.
            int_weights: See docstring of method :meth:`basic_rnn_step`.
            cm_weights: See docstring of method :meth:`basic_rnn_step`.
            ckpt_id: See docstring of method :meth:`basic_rnn_step`.
            is_last_step (bool): See docstring of method :meth:`basic_rnn_step`.

        Returns:
            (tuple): Tuple containing:

            - **h_t** (torch.Tensor): The tensor ``h_t`` of size
              ``[batch_size, n_hidden]`` with the new hidden state.
            - **c_t** (torch.Tensor): The tensor ``c_t`` of size
              ``[batch_size, n_hidden]`` with the new cell state.
        """
        use_cm = self._use_context_mod and d < self._num_rec_cm_layers
        # Determine the index of the hidden context mod layer.
        # Number of cm-layers per recurrent layer.
        n_cm_per_rec = self._context_mod_num_ts if \
            self._context_mod_num_ts != -1 and \
                self._context_mod_separate_layers_per_ts else 1
        cm_idx = self._cm_rnn_start_ind + d * n_cm_per_rec
        if self._context_mod_num_ts != -1 and \
                self._context_mod_separate_layers_per_ts:
            cm_idx += t

        c_t = h_t[1]
        h_t = h_t[0]
        HS = self._rnn_layers[d]

        if self._has_bias:
            assert len(int_weights) == 4
            weight_ih = int_weights[0]
            bias_ih = int_weights[1]
            weight_hh = int_weights[2]
            bias_hh = int_weights[3]
        else:
            assert len(int_weights) == 2
            weight_ih = int_weights[0]
            bias_ih = None
            weight_hh = int_weights[1]
            bias_hh = None

        # Compute total pre-activation input.
        gates = x_t @ weight_ih.t() + h_t @ weight_hh.t()
        if self.has_bias:
            gates += bias_ih + bias_hh

        i_t = gates[:, :HS]
        f_t = gates[:, HS:HS*2]
        g_t = gates[:, HS*2:HS*3]
        o_t = gates[:, HS*3:]

        # Compute activation.
        i_t = torch.sigmoid(i_t) # input
        f_t = torch.sigmoid(f_t) # forget
        g_t = torch.tanh(g_t)
        o_t = torch.sigmoid(o_t) # output

        # Compute c states.
        c_t = f_t * c_t + i_t * g_t

        # Note, we don't want to context-modulate the internal state c_t.
        # Otherwise, it might explode over timesteps since it wouldn't be
        # limited to [-1, 1] anymore. Instead, we only modulate the current
        # state (which is used to compute the current output h_t.
        c_t_mod = c_t

        # Context-dependent modulation (pre-activation).
        if use_cm and not self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                c_t_mod = self._context_mod_layers[cm_idx].forward(c_t_mod,
                    weights=cm_weights, ckpt_id=ckpt_id)

        # Compute h states.
        if self._a_fun is not None:
            h_t = o_t * self._a_fun(c_t_mod)
        else:
            h_t = o_t * c_t_mod

        # Context-dependent modulation (post-activation).
        # FIXME Christian: Shouldn't we apply the output gate `o_t` after
        # applying post-activation context-mod?
        if use_cm and self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                h_t = self._context_mod_layers[cm_idx].forward(h_t,
                    weights=cm_weights, ckpt_id=ckpt_id)

        return h_t, c_t

    def distillation_targets(self):
        """Targets to be distilled after training.

        See docstring of abstract super method
        :meth:`mnets.mnet_interface.MainNetInterface.distillation_targets`.

        This network does not have any distillation targets.

        Returns:
            ``None``
        """
        return None

    def get_output_weight_mask(self, out_inds=None, device=None):
        """Get masks to select output weights.

        See docstring of overwritten super method
        :meth:`mnets.mnet_interface.MainNetInterface.get_output_weight_mask`.
        """
        if len(self._fc_layers) > 0:
            return MainNetInterface.get_output_weight_mask(self,
                out_inds=out_inds, device=device)

        # TODO Output layer is recurrent. Hence, we have to properly handle
        # which weights contribute solely to certain output activations.
        raise NotImplementedError()

    def get_cm_weights(self):
        """Get internal maintained weights that are associated with context-
        modulation.

        Returns:
            (list): List of weights from
            :attr:`mnets.mnet_interface.MainNetInterface.internal_params` that
            are belonging to context-mod layers.
        """
        n_cm = self._num_context_mod_shapes()

        if n_cm == 0 or self._context_mod_no_weights:
            raise ValueError('Network maintains no context-modulation weights.')

        return self.internal_params[:n_cm]

    def get_non_cm_weights(self):
        """Get internal weights that are not associated with context-modulation.

        Returns:
            (list): List of weights from
            :attr:`mnets.mnet_interface.MainNetInterface.internal_params` that
            are not belonging to context-mod layers.
        """
        n_cm = 0 if self._context_mod_no_weights else \
            self._num_context_mod_shapes()

        return self.internal_params[n_cm:]

    def get_cm_inds(self):
        """Get the indices of
        :attr:`mnets.mnet_interface.MainNetInterface.param_shapes` that are
        associated with context-modulation.

        Returns:
            (list): List of integers representing indices of
            :attr:`mnets.mnet_interface.MainNetInterface.param_shapes`.
        """
        n_cm = self._num_context_mod_shapes()

        ret = []

        for i, meta in enumerate(self.param_shapes_meta):
            if meta['name'] == 'cm_shift' or meta['name'] == 'cm_scale':
                ret.append(i)
        assert n_cm == len(ret)

        return ret

    def init_hh_weights_orthogonal(self):
        """Initialize hidden-to-hidden weights orthogonally.

        This method will overwrite the hidden-to-hidden weights of recurrent
        layers.
        """
        for meta in self.param_shapes_meta:
            if meta['name'] == 'weight' and 'info' in meta.keys() and \
                    meta['info'] == 'hh' and meta['index'] != -1:
                print('Initializing hidden-to-hidden weights of recurrent ' +
                      'layer %d orthogonally.' % meta['layer'])
                W = self.internal_params[meta['index']]
                # LSTM weight matrices are stored such that the hidden-to-hidden 
                # matrices for the 4 gates are concatenated.
                if self.use_lstm:
                    out_dim, _ = W.shape
                    assert out_dim % 4 == 0
                    fs = out_dim // 4

                    W1 = W[:fs, :]
                    W2 = W[fs:2*fs, :]
                    W3 = W[2*fs:3*fs, :]
                    W4 = W[3*fs:, :]

                    torch.nn.init.orthogonal_(W1.data)
                    torch.nn.init.orthogonal_(W2.data)
                    torch.nn.init.orthogonal_(W3.data)
                    torch.nn.init.orthogonal_(W4.data)

                    # Sanity check to show that the init on partial matrices
                    # propagates back to the original tensor.
                    assert W[0,0] == W1[0,0]
                else:
                    torch.nn.init.orthogonal_(W.data)

    def _internal_weight_shapes(self):
        """Compute the tensor shapes of all internal weights (i.e., those not
        associated with context-modulation).

        Returns:
            (list): A list of list of integers, denoting the shapes of the
            individual parameter tensors.
        """
        coeff = 4 if self._use_lstm else 1
        shapes = []

        # Initial fully-connected layers.
        prev_dim = self._n_in
        for n_fc in self._fc_layers_pre:
            shapes.append([n_fc, prev_dim])
            if self._use_bias:
                shapes.append([n_fc])

            prev_dim = n_fc

        # Recurrent layers.
        for n_rec in self._rnn_layers:
            # Input-to-hidden
            shapes.append([n_rec*coeff, prev_dim])
            if self._use_bias:
                shapes.append([n_rec*coeff])

            # Hidden-to-hidden
            shapes.append([n_rec*coeff, n_rec])
            if self._use_bias:
                shapes.append([n_rec*coeff])

            if not self._use_lstm:
                # Hidden-to-output
                shapes.append([n_rec, n_rec])
                if self._use_bias:
                    shapes.append([n_rec])

            prev_dim = n_rec

        # Fully-connected layers.
        for n_fc in self._fc_layers:
            shapes.append([n_fc, prev_dim])
            if self._use_bias:
                shapes.append([n_fc])

            prev_dim = n_fc

        return shapes

if __name__ == '__main__':
    pass


#!/usr/bin/env python3
# Copyright 2019 <NAME>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# title          :mnets/simple_rnn.py
# author         :mc, be
# contact        :mariacer, <EMAIL>
# created        :10/28/2019
# version        :1.0
# python_version :3.6.8
"""
SimpleRNN
---------

Implementation of a simple recurrent neural network that has stacked vanilla RNN
or LSTM layers that are optionally enclosed by fully-connected layers.

An example usage is as a main model, where the main weights are initialized
and protected by a method such as EWC, and the context-modulation patterns of
the neurons are produced by an external hypernetwork.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np

from mnets.mnet_interface import MainNetInterface
from utils.torch_utils import init_params

class SimpleRNN(nn.Module, MainNetInterface):
    """Implementation of a simple RNN.

    This is a simple recurrent network, that receives input vector
    :math:`\mathbf{x}` and outputs a vector :math:`\mathbf{y}` of real values.

    Note:
        The output is non-linear if the last layer is recurrent! Otherwise,
        logits are returned (cmp. attribute
        :attr:`mnets.mnet_interface.MainNetInterface.has_fc_out`).

    Attributes:
        bptt_depth (int): The truncation depth for backprop through time.
            If ``-1``, backprop through time (BPTT) will unroll all timesteps
            present in the input. Otherwise, the forward pass will detach the
            RNN hidden states smaller or equal than
            ``num_timesteps - bptt_depth`` timesteps, resulting in truncated
            BPTT (T-BPTT).
        num_rec_layers (int): Number of recurrent layers in this network (i.e.,
            length of constructor argument ``rnn_layers``).
        use_lstm (bool): See constructor argument ``use_lstm``.

    Args:
        n_in (int): Number of inputs.
        rnn_layers (list or tuple): List of integers. Each entry denotes the
            size of a recurrent layer. Recurrent layers will simply be stacked
            as layers of this network.

            If ``fc_layers_pre`` is empty, then the recurrent layers are the
            initial layers.
            If ``fc_layers`` is empty, then the last entry of this list will
            denote the output size.

            Note:
                This list may never be empty.
        fc_layers_pre (list or tuple): List of integers. Before the recurrent
            layers a set of fully-connected layers may be added. This might be
            specially useful when constructing recurrent autoencoders. The
            entries of this list will denote the sizes of those layers.

            If ``fc_layers_pre`` is not empty, its first entry will denote the
            input size of this network.
        fc_layers (list or tuple): List of integers. After the recurrent layers,
            a set of fully-connected layers is added. The entries of this list
            will denote the sizes of those layers.

            If ``fc_layers`` is not empty, its last entry will denote the output
            size of this network.
        activation: The nonlinearity used in hidden layers.
        use_lstm (bool): If set to `True``, the recurrent layers will be LSTM
            layers.
        use_bias (bool): Whether layers may have bias terms.
        no_weights (bool): If set to ``True``, no trainable parameters will be
            constructed, i.e., weights are assumed to be produced ad-hoc
            by a hypernetwork and passed to the :meth:`forward` method.
        init_weights (list, optional): This option is for convinience reasons.
            The option expects a list of parameter values that are used to
            initialize the network weights. As such, it provides a
            convinient way of initializing a network with a weight draw
            produced by the hypernetwork.

            Note, internal weights (see 
            :attr:`mnets.mnet_interface.MainNetInterface.weights`) will be
            affected by this argument only.
        kaiming_rnn_init (bool): By default, PyTorch initializes its recurrent
            layers uniformly with an interval defined by the square-root of the
            inverse of the layer size.

            If this option is enabled, then the recurrent layers will be
            initialized using the kaiming init as implemented by the function
            :func:`utils.torch_utils.init_params`.
        context_mod_last_step (bool): Whether context modulation is applied
            at the last time step os a recurrent layer only. If ``False``,
            context modulation is applied at every time step.

            Note:
                This option only applies if ``use_context_mod`` is ``True``.
        context_mod_num_ts (int, optional): The maximum number of timesteps.
            If specified, context-modulation with a different set of weights is
            applied at every timestep. If ``context_mod_separate_layers_per_ts``
            is ``True``, then a separate context-mod layer per timestep will be
            created. Otherwise, a single context-mod layer is created, but the
            expected parameter shapes for this layer are
            ``[context_mod_num_ts, *context_mod_shape]``.

            Note:
                This option only applies if ``use_context_mod`` is ``True``.
        context_mod_separate_layers_per_ts (bool): If specified, a separate
            context-mod layer per timestep is created (required if
            ``context_mod_no_weights`` is ``False``).
            
            Note:
                Only applies if ``context_mod_num_ts`` is specified.
        verbose (bool): Whether to print information (e.g., the number of
            weights) during the construction of the network.
        **kwargs: Keyword arguments regarding context modulation. This class
            can process the same context-modulation related arguments as class
            :class:`mnets.mlp.MLP` (plus the additional ones noted above).
    """
    def __init__(self, n_in=1, rnn_layers=(10,), fc_layers_pre=(),
                 fc_layers=(1,), activation=torch.nn.Tanh(), use_lstm=False,
                 use_bias=True, no_weights=False,
                 init_weights=None, kaiming_rnn_init=False,
                 context_mod_last_step=False,
                 context_mod_num_ts=-1,
                 context_mod_separate_layers_per_ts=False,
                 verbose=True,
                 **kwargs):
        # FIXME find a way using super to handle multiple inheritance.
        nn.Module.__init__(self)
        MainNetInterface.__init__(self)

        self._bptt_depth = -1

        # FIXME Arbitrary restriction.
        if activation is None or isinstance(activation, (torch.nn.ReLU, \
                torch.nn.Tanh)):
            self._a_fun = activation
        else:
            raise ValueError('Only linear, relu and tanh activations are ' + \
                             'allowed for recurrent networks.')

        if len(rnn_layers) == 0:
            raise ValueError('The network always needs to have at least one ' +
                             'recurrent layer.')
        if len(fc_layers) == 0:
            has_rec_out_layer = True
            #n_out = rnn_layers[-1]
        else:
            has_rec_out_layer = False
            #n_out = fc_layers[-1]

        self._n_in = n_in
        self._rnn_layers = rnn_layers
        self._fc_layers_pre = fc_layers_pre
        self._fc_layers = fc_layers

        self._no_weights = no_weights

        ### Parse or set context-mod arguments ###
        rem_kwargs = MainNetInterface._parse_context_mod_args(kwargs)
        if len(rem_kwargs) > 0:
            raise ValueError('Keyword arguments %s unknown.' % str(rem_kwargs))

        self._use_context_mod = kwargs['use_context_mod']
        self._context_mod_inputs = kwargs['context_mod_inputs']
        self._no_last_layer_context_mod = kwargs['no_last_layer_context_mod']
        self._context_mod_no_weights = kwargs['context_mod_no_weights']
        self._context_mod_post_activation = \
            kwargs['context_mod_post_activation']
        self._context_mod_gain_offset = kwargs['context_mod_gain_offset']
        self._context_mod_gain_softplus = kwargs['context_mod_gain_softplus']

        # Context-mod options specific to RNNs
        self._context_mod_last_step = context_mod_last_step
        # FIXME We have to specify this option even if
        # `context_mod_separate_layers_per_ts` is False (in order to set
        # sensible parameter shapes). However, the forward method can deal with
        # an arbitrary timestep length.
        self._context_mod_num_ts = context_mod_num_ts
        self._context_mod_separate_layers_per_ts = \
            context_mod_separate_layers_per_ts

        # More appropriate naming of option.
        self._context_mod_outputs = not self._no_last_layer_context_mod

        if context_mod_num_ts != -1:
            if context_mod_last_step:
                raise ValueError('Options "context_mod_last_step" and ' +
                                 '"context_mod_num_ts" are not compatible.')
            if not self._context_mod_no_weights and \
                    not context_mod_separate_layers_per_ts:
                raise ValueError('When applying context-mod per timestep ' +
                    'while maintaining weights internally, option' +
                    '"context_mod_separate_layers_per_ts" must be set.')
        ### Parse or set context-mod arguments - DONE ###

        self._has_bias = use_bias
        self._has_fc_out = True if not has_rec_out_layer else False
        # We need to make sure that the last 2 entries of `weights` correspond
        # to the weight matrix and bias vector of the last layer.
        self._mask_fc_out = True if not has_rec_out_layer else False
        # Note, recurrent layers always use non-linearities and their activities
        # are squashed by a non-linearity (otherwise, internal states could
        # vanish/explode with increasing sequence length).
        self._has_linear_out = True if not has_rec_out_layer else False

        self._param_shapes = []
        self._param_shapes_meta = []

        self._weights = None if no_weights and self._context_mod_no_weights \
            else nn.ParameterList()
        self._hyper_shapes_learned = None \
            if not no_weights and not self._context_mod_no_weights else []
        self._hyper_shapes_learned_ref = None if self._hyper_shapes_learned \
            is None else []

        self._layer_weight_tensors = nn.ParameterList()
        self._layer_bias_vectors = nn.ParameterList()

        self._use_lstm = use_lstm
        if use_lstm:
            self._rnn_fct = self.lstm_rnn_step
        else:
            self._rnn_fct = self.basic_rnn_step

        #################################################
        ### Define and initialize context mod weights ###
        #################################################

        # The context-mod layers consist of sequential layers ordered as:
        # - initial fully-connected layers if len(fc_layers_pre)>0
        # - recurrent layers
        # - final fully-connected layers if len(fc_layers)>0

        self._context_mod_layers = nn.ModuleList() if self._use_context_mod \
            else None

        self._cm_rnn_start_ind = 0
        self._num_fc_cm_layers = None
        if self._use_context_mod:
            cm_layer_inds = []
            cm_shapes = []

            # Gather sizes of all activation vectors within the network that
            # will be subject to context-modulation.
            if self._context_mod_inputs:
                self._cm_rnn_start_ind += 1
                cm_shapes.append([n_in])

                # We reserve layer zero for input context-mod. Otherwise, there
                # is no layer zero.
                cm_layer_inds.append(0)

            if len(fc_layers_pre) > 0:
                self._cm_rnn_start_ind += len(fc_layers_pre)

            # We use odd numbers for actual layers and even number for all
            # context-mod layers.
            rem_cm_inds = range(2, 2*(len(fc_layers_pre)+len(rnn_layers)+\
                len(fc_layers))+1, 2)

            num_rec_cm_layers = len(rnn_layers)
            if has_rec_out_layer and not self._context_mod_outputs:
                num_rec_cm_layers -= 1
            self._num_rec_cm_layers = num_rec_cm_layers

            jj = 0
            # Add initial fully-connected context-mod layers.
            num_fc_pre_cm_layers = len(fc_layers_pre)
            self._num_fc_pre_cm_layers = num_fc_pre_cm_layers
            for i in range(num_fc_pre_cm_layers):
                cm_shapes.append([fc_layers_pre[i]])
                cm_layer_inds.append(rem_cm_inds[jj])
                jj += 1

            # Add recurrent context-mod layers.
            for i in range(num_rec_cm_layers):
                if context_mod_num_ts != -1:
                    if context_mod_separate_layers_per_ts:
                        cm_rnn_shapes = [[rnn_layers[i]]] * context_mod_num_ts
                    else:
                        # Only a single context-mod layer will be added, but we
                        # directly edit the correponding `param_shape` later.
                        assert self._context_mod_no_weights
                        cm_rnn_shapes = [[rnn_layers[i]]]
                else:
                    cm_rnn_shapes = [[rnn_layers[i]]]

                cm_shapes.extend(cm_rnn_shapes)
                cm_layer_inds.extend([rem_cm_inds[jj]] * len(cm_rnn_shapes))
                jj += 1

            # Add final fully-connected context-mod layers.
            num_fc_cm_layers = len(fc_layers)
            if num_fc_cm_layers > 0 and not self._context_mod_outputs:
                num_fc_cm_layers -= 1
            self._num_fc_cm_layers = num_fc_cm_layers
            for i in range(num_fc_cm_layers):
                cm_shapes.append([fc_layers[i]])
                cm_layer_inds.append(rem_cm_inds[jj])
                jj += 1

            self._add_context_mod_layers(cm_shapes, cm_layers=cm_layer_inds)

            if context_mod_num_ts != -1 and not \
                    context_mod_separate_layers_per_ts:
                # In this case, there is only one context-mod layer for each
                # recurrent layer, but we want to have separate weights per
                # timestep.
                # Hence, we adapt the expected parameter shape, such that we
                # get a different set of weights per timestep. This will be
                # split into multiple weights that are succesively fed into the
                # same layer inside the forward method.

                for i in range(num_rec_cm_layers):
                    cmod_layer = \
                        self.context_mod_layers[self._cm_rnn_start_ind+i]
                    cm_shapes_rnn = [[context_mod_num_ts, *s] for s in \
                                      cmod_layer.param_shapes]

                    ps_ind = int(np.sum([ \
                        len(self.context_mod_layers[ii].param_shapes) \
                        for ii in range(self._cm_rnn_start_ind+i)]))
                    self._param_shapes[ps_ind:ps_ind+len(cm_shapes_rnn)] = \
                        cm_shapes_rnn
                    assert self._hyper_shapes_learned is not None
                    self._hyper_shapes_learned[ \
                        ps_ind:ps_ind+len(cm_shapes_rnn)] = cm_shapes_rnn

        ########################
        ### Internal weights ###
        ########################
        prev_dim = self._n_in

        def define_fc_layer_weights(fc_layers, prev_dim, num_prev_layers):
            """Define the weights and shapes of the fully-connected layers.

            Args:
                fc_layers (list): The list of fully-connected layer dimensions.
                prev_dim (int): The output size of the previous layer.
                num_prev_layers (int): The number of upstream layers to the 
                    current one (a layer with its corresponding
                    context-mod layer(s) count as one layer). Count should
                    start at ``1``.

            Returns:
                (int): The output size of the last fully-connected layer
                considered here.
            """
            # FIXME We should instead build an MLP instance. But then we still
            # have to adapt all attributes accordingly.
            for i, n_fc in enumerate(fc_layers):
                s_w = [n_fc, prev_dim]
                s_b = [n_fc] if self._has_bias else None

                for j, s in enumerate([s_w, s_b]):
                    if s is None:
                        continue

                    is_bias = True
                    if j % 2 == 0:
                        is_bias = False

                    if not self._no_weights:
                        self._weights.append(nn.Parameter(torch.Tensor(*s),
                                                          requires_grad=True))
                        if is_bias:
                            self._layer_bias_vectors.append(self._weights[-1])
                        else:
                            self._layer_weight_tensors.append(self._weights[-1])
                    else:
                        self._hyper_shapes_learned.append(s)
                        self._hyper_shapes_learned_ref.append( \
                            len(self.param_shapes))

                    self._param_shapes.append(s)
                    self._param_shapes_meta.append({
                        'name': 'bias' if is_bias else 'weight',
                        'index': -1 if self._no_weights else \
                            len(self._weights)-1,
                        'layer': i * 2 + num_prev_layers, # Odd numbers
                    })

                prev_dim = n_fc

            return prev_dim

        ### Initial fully-connected layers.
        prev_dim = define_fc_layer_weights(self._fc_layers_pre, prev_dim, 1)

        ### Recurrent layers.
        coeff = 4 if self._use_lstm else 1
        for i, n_rec in enumerate(self._rnn_layers):
            # Input-to-hidden
            s_w_ih = [n_rec*coeff, prev_dim]
            s_b_ih = [n_rec*coeff] if use_bias else None

            # Hidden-to-hidden
            s_w_hh = [n_rec*coeff, n_rec]
            s_b_hh = [n_rec*coeff] if use_bias else None

            # Hidden-to-output.
            # Note, for an LSTM cell, the hidden state vector is also the
            # output vector.
            if not self._use_lstm:
                s_w_ho = [n_rec, n_rec]
                s_b_ho = [n_rec] if use_bias else None
            else:
                s_w_ho = None
                s_b_ho = None

            for j, s in enumerate([s_w_ih, s_b_ih, s_w_hh, s_b_hh, s_w_ho,
                                   s_b_ho]):
                if s is None:
                    continue

                is_bias = True
                if j % 2 == 0:
                    is_bias = False

                wtype = 'ih'
                if 2 <= j < 4:
                    wtype = 'hh'
                elif j >=4:
                    wtype = 'ho'

                if not no_weights:
                    self._weights.append(nn.Parameter(torch.Tensor(*s),
                                                      requires_grad=True))
                    if is_bias:
                        self._layer_bias_vectors.append(self._weights[-1])
                    else:
                        self._layer_weight_tensors.append(self._weights[-1])
                else:
                    self._hyper_shapes_learned.append(s)
                    self._hyper_shapes_learned_ref.append( \
                        len(self.param_shapes))

                self._param_shapes.append(s)
                self._param_shapes_meta.append({
                    'name': 'bias' if is_bias else 'weight',
                    'index': -1 if no_weights else len(self._weights)-1,
                    'layer': i * 2 + 1 + 2 * len(fc_layers_pre), # Odd numbers
                    'info': wtype
                })

            prev_dim = n_rec

        ### Fully-connected layers.
        prev_dim = define_fc_layer_weights(self._fc_layers, prev_dim, \
            1 + 2 * len(fc_layers_pre) + 2 * len(rnn_layers))

        ### Initialize weights.
        if init_weights is not None:
            assert self._weights is not None
            assert len(init_weights) == len(self.weights)
            for i in range(len(init_weights)):
                assert np.all(np.equal(list(init_weights[i].shape),
                                       self.weights[i].shape))
                self.weights[i].data = init_weights[i]
        else:
            rec_start = len(fc_layers_pre)
            rec_end = rec_start + len(rnn_layers) * (2 if use_lstm else 3)
            # Note, Pytorch applies a uniform init to its recurrent layers, as
            # defined here:
            # https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py#L155
            for i in range(len(self._layer_weight_tensors)):
                if i >=rec_start and i < rec_end:
                    # Recurrent layer weights.
                    if kaiming_rnn_init:
                        init_params(self._layer_weight_tensors[i],
                            self._layer_bias_vectors[i] if use_bias else None)
                    else:
                        a = 1.0 / math.sqrt(rnn_layers[(i-rec_start) // \
                            (2 if use_lstm else 3)])
                        nn.init.uniform_(self._layer_weight_tensors[i], -a, a)
                        if use_bias:
                            nn.init.uniform_(self._layer_bias_vectors[i], -a, a)
                else:
                    # FC layer weights.
                    init_params(self._layer_weight_tensors[i],
                        self._layer_bias_vectors[i] if use_bias else None)

        num_weights = MainNetInterface.shapes_to_num_weights(self._param_shapes)

        if verbose:
            if self._use_context_mod:
                cm_num_weights =  \
                    MainNetInterface.shapes_to_num_weights(cm_shapes)

            print('Creating a simple RNN with %d weights' % num_weights
                  + (' (including %d weights associated with-' % cm_num_weights
                     + 'context modulation)' if self._use_context_mod else '')
                  + '.')

        self._is_properly_setup()

    @property
    def bptt_depth(self):
        """Getter for attribute :attr:`bptt_depth`."""
        return self._bptt_depth

    @bptt_depth.setter
    def bptt_depth(self, value):
        """Setter for attribute :attr:`bptt_depth`."""
        self._bptt_depth = value

    @property
    def num_rec_layers(self):
        """Getter for read-only attribute :attr:`num_rec_layers`."""
        return len(self._rnn_layers)

    @property
    def use_lstm(self):
        """Getter for read-only attribute :attr:`use_lstm`."""
        return self._use_lstm

    def split_cm_weights(self, cm_weights, condition, num_ts=0):
        """Split context-mod weights per context-mod layer.

        Args:
            cm_weights (torch.Tensor): All context modulation weights.
            condition (optional, int): If provided, then this argument will be
                passed as argument ``ckpt_id`` to the method
                :meth:`utils.context_mod_layer.ContextModLayer.forward`.
            num_ts (int): The length of the sequences.

        Returns:
            (Tuple): Where the tuple contains:

            - **cm_inputs_weights**: The cm input weights.
            - **cm_fc_pre_layer_weights**: The cm pre-recurrent weights.
            - **cm_rec_layer_weights**: The cm recurrent weights.
            - **cm_fc_layer_weights**: The cm post-recurrent weights.
            - **n_cm_rec**: The number of recurrent cm layers.
            - **cmod_cond**: The context-mod condition.
        """

        n_cm_rec = -1
        cm_fc_pre_layer_weights = None
        cm_fc_layer_weights = None
        cm_inputs_weights = None
        cm_rec_layer_weights = None
        if cm_weights is not None:
            if self._context_mod_num_ts != -1 and \
                    self._context_mod_separate_layers_per_ts:
                assert num_ts <= self._context_mod_num_ts

            # Note, an mnet layer might contain multiple context-mod layers
            # (a recurrent layer can have a separate context-mod layer per
            # timestep).
            cm_fc_pre_layer_weights = []
            cm_rec_layer_weights = [[] for _ in range(self._num_rec_cm_layers)]
            cm_fc_layer_weights = []

            # Number of cm-layers per recurrent layer.
            n_cm_per_rec = self._context_mod_num_ts if \
                self._context_mod_num_ts != -1 and \
                    self._context_mod_separate_layers_per_ts else 1
            n_cm_rec = n_cm_per_rec * self._num_rec_cm_layers

            cm_start = 0
            for i, cm_layer in enumerate(self.context_mod_layers):
                cm_end = cm_start + len(cm_layer.param_shapes)

                if i == 0 and self._context_mod_inputs:
                    cm_inputs_weights = cm_weights[cm_start:cm_end]
                elif i < self._cm_rnn_start_ind:
                    cm_fc_pre_layer_weights.append(cm_weights[cm_start:cm_end])
                elif i >= self._cm_rnn_start_ind and \
                        i < self._cm_rnn_start_ind + n_cm_rec:
                    # Index of recurrent layer.
                    i_r = (i-self._cm_rnn_start_ind) // n_cm_per_rec
                    cm_rec_layer_weights[i_r].append( \
                        cm_weights[cm_start:cm_end])
                else:
                    cm_fc_layer_weights.append(cm_weights[cm_start:cm_end])
                cm_start = cm_end

            # We need to split the context-mod weights in the following case,
            # as they are currently just stacked on top of each other.
            if self._context_mod_num_ts != -1 and \
                    not self._context_mod_separate_layers_per_ts:
                for i, cm_w_list in enumerate(cm_rec_layer_weights):
                    assert len(cm_w_list) == 1

                    cm_rnn_weights = cm_w_list[0]
                    cm_rnn_layer = self.context_mod_layers[ \
                        self._cm_rnn_start_ind+i]

                    assert len(cm_rnn_weights) == len(cm_rnn_layer.param_shapes)
                    # The first dimension are the weights of this layer per
                    # timestep.
                    num_ts_cm = -1
                    for j, s in enumerate(cm_rnn_layer.param_shapes):
                        assert len(cm_rnn_weights[j].shape) == len(s) + 1
                        if j == 0:
                            num_ts_cm = cm_rnn_weights[j].shape[0]
                        else:
                            assert num_ts_cm == cm_rnn_weights[j].shape[0]
                    assert num_ts <= num_ts_cm

                    cm_w_chunked = [None] * len(cm_rnn_weights)
                    for j, cm_w in enumerate(cm_rnn_weights):
                        cm_w_chunked[j] = torch.chunk(cm_w, num_ts_cm, dim=0)

                    # Now we gather all these chunks to assemble the weights
                    # needed per timestep (as if
                    # `_context_mod_separate_layers_per_t` were True).
                    cm_w_list = []
                    for j in range(num_ts_cm):
                        tmp_list = []
                        for chunk in cm_w_chunked:
                            tmp_list.append(chunk[j].squeeze(dim=0))
                        cm_w_list.append(tmp_list)
                    cm_rec_layer_weights[i] = cm_w_list

            # Note, the last layer does not necessarily have context-mod
            # (depending on `self._context_mod_outputs`).
            if len(cm_rec_layer_weights) < len(self._rnn_layers):
                cm_rec_layer_weights.append(None)
            if len(cm_fc_layer_weights) < len(self._fc_layers):
                cm_fc_layer_weights.append(None)


        #######################
        ### Parse condition ###
        #######################
        cmod_cond = None
        if condition is not None:
            assert isinstance(condition, int)
            cmod_cond = condition

            # Note, the cm layer will ignore the cmod condition if weights
            # are passed.
            # FIXME Find a more elegant solution.
            cm_inputs_weights = None
            cm_fc_pre_layer_weights = [None] * len(cm_fc_pre_layer_weights)
            cm_rec_layer_weights = [[None] * len(cm_ws) for cm_ws in \
                                    cm_rec_layer_weights]
            cm_fc_layer_weights = [None] * len(cm_fc_layer_weights)

        return cm_inputs_weights, cm_fc_pre_layer_weights, cm_fc_layer_weights,\
            cm_rec_layer_weights, n_cm_rec, cmod_cond

    def split_internal_weights(self, int_weights):
        """Split internal weights per layer.

        Args:
            int_weights (torch.Tensor): All internal weights.

        Returns:
            (Tuple): Where the tuple contains:

            - **fc_pre_w_weights**: The pre-recurrent w weights.
            - **fc_pre_b_weights**: The pre-recurrent b weights.
            - **rec_weights**: The recurrent weights.
            - **fc_w_weights**:The post-recurrent w weights.
            - **fc_b_weights**: The post-recurrent b weights.
        """
        n_cm = self._num_context_mod_shapes()

        int_meta = self.param_shapes_meta[n_cm:]
        assert len(int_meta) == len(int_weights)
        fc_pre_w_weights = []
        fc_pre_b_weights = []
        rec_weights =[[] for _ in range(len(self._rnn_layers))]
        fc_w_weights = []
        fc_b_weights = []

        # Number of pre-fc weights in total.
        n_fc_pre = len(self._fc_layers_pre)
        if self.has_bias:
            n_fc_pre *= 2

        # Number of weights per recurrent layer.
        if self._use_lstm:
            n_rw = 4 if self.has_bias else 2
        else:
            n_rw = 6 if self.has_bias else 3

        for i, w in enumerate(int_weights):
            if i < n_fc_pre: # fc pre weights
                if int_meta[i]['name'] == 'weight':
                    fc_pre_w_weights.append(w)
                else:
                    assert int_meta[i]['name'] == 'bias'
                    fc_pre_b_weights.append(w)
            elif i >= n_fc_pre and \
                    i < n_rw * len(self._rnn_layers) + n_fc_pre: # recurrent w
                r_ind = (i - n_fc_pre) // n_rw
                rec_weights[r_ind].append(w)
            else: # fc weights
                if int_meta[i]['name'] == 'weight':
                    fc_w_weights.append(w)
                else:
                    assert int_meta[i]['name'] == 'bias'
                    fc_b_weights.append(w)

        if not self.has_bias:
            assert len(fc_pre_b_weights) == 0
            fc_pre_b_weights = [None] * len(fc_pre_w_weights)

            assert len(fc_b_weights) == 0
            fc_b_weights = [None] * len(fc_w_weights)

        return fc_pre_w_weights, fc_pre_b_weights, rec_weights, fc_w_weights, \
            fc_b_weights

    def split_weights(self, weights):
        """Split weights into internal and context-mod weights.

        Extract which weights should be used,  I.e., are we using internally
        maintained weights or externally given ones or are we even mixing
        between these groups.

        Args:
            weights (torch.Tensor): All weights.

        Returns:
            (Tuple): Where the tuple contains:

            - **int_weights**: The internal weights.
            - **cm_weights**: The context-mod weights.
        """
        n_cm = self._num_context_mod_shapes()

        ### FIXME Code copied from MLP its `forward` method ###

        # Make sure cm_weights are either `None` or have the correct dimensions.
        if weights is None:
            weights = self.weights

            if self._use_context_mod:
                cm_weights = weights[:n_cm]
                int_weights = weights[n_cm:]
            else:
                cm_weights = None
                int_weights = weights
        else:
            int_weights = None
            cm_weights = None

            if isinstance(weights, dict):
                assert 'internal_weights' in weights.keys() or \
                       'mod_weights' in weights.keys()
                if 'internal_weights' in weights.keys():
                    int_weights = weights['internal_weights']
                if 'mod_weights' in weights.keys():
                    cm_weights = weights['mod_weights']
            else:
                if self._use_context_mod and \
                        len(weights) == n_cm:
                    cm_weights = weights
                else:
                    assert len(weights) == len(self.param_shapes)
                    if self._use_context_mod:
                        cm_weights = weights[:n_cm]
                        int_weights = weights[n_cm:]
                    else:
                        int_weights = weights

            if self._use_context_mod and cm_weights is None:
                if self._context_mod_no_weights:
                    raise Exception('Network was generated without weights ' +
                        'for context-mod layers. Hence, they must be passed ' +
                        'via the "weights" option.')
                cm_weights = self.weights[:n_cm]
            if int_weights is None:
                if self._no_weights:
                    raise Exception('Network was generated without internal ' +
                        'weights. Hence, they must be passed via the ' +
                        '"weights" option.')
                if self._context_mod_no_weights:
                    int_weights = self.weights
                else:
                    int_weights = self.weights[n_cm:]

            # Note, context-mod weights might have different shapes, as they
            # may be parametrized on a per-sample basis.
            if self._use_context_mod:
                assert len(cm_weights) == n_cm
            int_shapes = self.param_shapes[n_cm:]
            assert len(int_weights) == len(int_shapes)
            for i, s in enumerate(int_shapes):
                assert np.all(np.equal(s, list(int_weights[i].shape)))

        ### FIXME Code copied until here ###
        return int_weights, cm_weights

    def forward(self, x, weights=None, distilled_params=None, condition=None,
                return_hidden=False, return_hidden_int=False):
        """Compute the output :math:`y` of this network given the input
        :math:`x`.

        Args:
            (....): See docstring of method
                :meth:`mnets.mnet_interface.MainNetInterface.forward`. We
                provide some more specific information below.
            weights (list or dict): See argument ``weights`` of method
                :meth:`mnets.mlp.MLP.forward`.
            condition (optional, int): If provided, then this argument will be
                passed as argument ``ckpt_id`` to the method
                :meth:`utils.context_mod_layer.ContextModLayer.forward`.
            return_hidden (bool, optional): If ``True``, all hidden activations
                of fully-connected and recurrent layers (where we defined
                :math:`y_t` as hidden state of vannila RNN layers as these are
                the layer outputs passed to the next layer) are returned.
                recurrent hidden activations will be returned.
            return_hidden_int (bool, optional): If ``True``, in addition to
                ``hidden``, an additional variable ``hidden_int`` is returned
                containing the internal hidden states of recurrent layers (i.e.,
                the cell states :math:`c_t` for LSTMs and the actual hidden
                state :math:`h_t` for Elman layers) are returned.

        Returns:
            (torch.Tensor or tuple): Where the tuple is containing:

            - **output** (torch.Tensor): The output of the network.
            - **hidden** (list): If ``return_hidden`` is ``True``, then the
              hidden activities of each layer are returned, which have the shape
              ``(seq_length, batch_size, n_hidden)``.
            - **hidden_int**: If ``return_hidden_int`` is ``True``, then in
              addition to ``hidden`` a tensor ``hidden_int`` is returned
              containing internal hidden states of recurrent layers.
        """
        assert distilled_params is None

        if ((not self._use_context_mod and self._no_weights) or \
                (self._no_weights or self._context_mod_no_weights)) and \
                weights is None:
            raise Exception('Network was generated without weights. ' +
                            'Hence, "weights" option may not be None.')

        #######################
        ### Extract weights ###
        #######################
        # Extract which weights should be used.
        int_weights, cm_weights = self.split_weights(weights)

        ### Split context-mod weights per context-mod layer.
        cm_inputs_weights, cm_fc_pre_layer_weights, cm_fc_layer_weights, \
            cm_rec_layer_weights, n_cm_rec, cmod_cond = self.split_cm_weights(
                cm_weights, condition, num_ts=x.shape[0])

        ### Extract internal weights.
        fc_pre_w_weights, fc_pre_b_weights, rec_weights, fc_w_weights, \
            fc_b_weights = self.split_internal_weights(int_weights)

        ###########################
        ### Forward Computation ###
        ###########################
        ret_hidden = None
        if return_hidden:
            ret_hidden = []

        h = x

        cm_offset = 0
        if self._use_context_mod and self._context_mod_inputs:
            cm_offset += 1
            # Apply context modulation in the inputs.
            h = self._context_mod_layers[0].forward(h,
                weights=cm_inputs_weights, ckpt_id=cmod_cond, bs_dim=1)

        ### Initial fully-connected layer activities.
        ret_hidden, h = self.compute_fc_outputs(h, fc_pre_w_weights, \
            fc_pre_b_weights, len(self._fc_layers_pre), \
            cm_fc_pre_layer_weights, cm_offset, cmod_cond, False, ret_hidden)

        ### Recurrent layer activities.
        ret_hidden_int = [] # the internal hidden activations
        for d in range(len(self._rnn_layers)):
            if self._use_context_mod:
                h, h_int = self.compute_hidden_states(h, d, rec_weights[d],
                    cm_rec_layer_weights[d], cmod_cond)
            else:
                h, h_int = self.compute_hidden_states(h, d, rec_weights[d],
                                                      None, None)
            if ret_hidden is not None:
                ret_hidden.append(h)
                ret_hidden_int.append(h_int)

        ### Fully-connected layer activities.
        cm_offset = self._cm_rnn_start_ind + n_cm_rec
        ret_hidden, h = self.compute_fc_outputs(h, fc_w_weights, fc_b_weights, \
            self._num_fc_cm_layers, cm_fc_layer_weights, cm_offset, cmod_cond,
            True, ret_hidden)

        # FIXME quite ugly
        if return_hidden:
            # The last element is the output activity.
            ret_hidden.pop()
            if return_hidden_int:
                return h, ret_hidden, ret_hidden_int
            else:
                return h, ret_hidden
        else:
            return h

    def compute_fc_outputs(self, h, fc_w_weights, fc_b_weights, num_fc_cm_layers,
                cm_fc_layer_weights, cm_offset, cmod_cond, is_post_fc, 
                ret_hidden):
        """Compute the forward pass through the fully-connected layers.

        This method also appends activations to ``ret_hidden``.

        Args:
            h (torch.Tensor): The input from the previous layer.
            fc_w_weights (list): The weights for the fc layers.
            fc_b_weights (list): The biases for the fc layers.
            num_fc_cm_layers (int): The number of context-modulation
                layers associated with this set of fully-connected layers.
            cm_fc_layer_weights (list): The context-modulation weights
                associated with the current layers.
            cm_offset (int): The index to access the correct context-mod
                layers.
            cmod_cond (bool): Some condition to perform context modulation.
            is_post_fc (bool); Whether those layers are applied as last
                layers of the network. In this case, there will be no
                activation applied to the last layer outputs.
            ret_hidden (list or None): The hidden recurrent activations.

        Return:
            (Tuple): Tuple containing:

            - **ret_hidden**: The hidden recurrent activations.
            - **h**: Transformed activation ``h``.
        """
        for d in range(len(fc_w_weights)):
            use_cm = self._use_context_mod and d < num_fc_cm_layers
            # Compute output.
            h = F.linear(h, fc_w_weights[d], bias=fc_b_weights[d])

            # Context-dependent modulation (pre-activation).
            if use_cm and not self._context_mod_post_activation:
                h = self._context_mod_layers[cm_offset+d].forward(h,
                    weights=cm_fc_layer_weights[d], ckpt_id=cmod_cond, 
                    bs_dim=1)

            # Non-linearity
            # Note, non-linearity is not applied to outputs of the network.
            if self._a_fun is not None and \
                    (not is_post_fc or d < len(fc_w_weights)-1):
                h = self._a_fun(h)

            # Context-dependent modulation (post-activation).
            if use_cm and self._context_mod_post_activation:
                h = self._context_mod_layers[cm_offset+d].forward(h,
                    weights=cm_fc_layer_weights[d], ckpt_id=cmod_cond, 
                    bs_dim=1)

            if ret_hidden is not None:
                ret_hidden.append(h)

        return ret_hidden, h

    def compute_hidden_states(self, x, layer_ind, int_weights, cm_weights,
                              ckpt_id, h_0=None, c_0=None):
        """Compute the hidden states for the recurrent layer ``layer_ind`` from
        a sequence of inputs :math:`x`.

        If so specified, context modulation is applied before or after the
        nonlinearities.

        Args:
            x: The inputs :math:`x` to the layer. :math:`x` has shape
                ``[sequence_len, batch_size, n_hidden_prev]``.
            layer_ind (int): Index of the layer.
            int_weights: Internal weights associated with this recurrent layer.
            cm_weights: Context modulation weights.
            ckpt_id: Will be passed as option ``ckpt_id`` to method
                :meth:`utils.context_mod_layer.ContextModLayer.forward` if
                context-mod layers are used.
            h_0 (torch.Tensor, optional): The initial state for :math:`h`.
            c_0 (torch.Tensor, optional): The initial state for :math:`c`. Note
                that for LSTMs, if the initial state is to be defined, this
                variable is necessary also, not only :math:`h_0`, whereas for
                vanilla RNNs it is enough to provide :math:`h_0` as :math:`c_0`
                represents the output of the layer and it can be easily computed
                from `h_0`.

        Returns:
            (tuple): Tuple containing:

            - **outputs** (torch.Tensor): The sequence of visible hidden states
              given the input. It has shape 
              ``[sequence_len, batch_size, n_hidden]``.
            - **hiddens** (torch.Tensor): The sequence of hidden states given
              the input. For LSTMs, this corresponds to :math:`c`.
              It has shape ``[sequence_len, batch_size, n_hidden]``.
        """
        seq_length, batch_size, n_hidden_prev = x.shape
        n_hidden = self._rnn_layers[layer_ind]

        # Generate initial hidden states.
        # Note that for LSTMs h_0 is the hidden state and output vector whereas
        # c_0 is the internal cell state vector.
        # For a vanilla RNN h_0 is the hidden state whereas c_0 is the output
        # vector.
        if h_0 is None:
            h_0 = (torch.zeros(batch_size, n_hidden, device=x.device))
        if c_0 is None:
            c_0 = (torch.zeros(batch_size, n_hidden, device=x.device))
        assert h_0.shape[0] == c_0.shape[0] == batch_size
        assert h_0.shape[1] == c_0.shape[1] == n_hidden

        # If we want to apply context modulation in each time step, we need
        # to split the input sequence and call pytorch function at every
        # time step.
        outputs = []
        hiddens = []
        h_t = h_0
        c_t = c_0
        for t in range(seq_length):
            x_t = x[t,:,:]

            if cm_weights is not None and self._context_mod_num_ts != -1:
                curr_cm_weights = cm_weights[t]
            elif cm_weights is not None:
                assert len(cm_weights) == 1
                curr_cm_weights = cm_weights[0]
            else:
                curr_cm_weights = cm_weights

            # Compute the actual rnn step (either vanilla or LSTM, depending on
            # the flag self._use_lstm).
            is_last_step = t==(seq_length-1)
            h_t, c_t = self._rnn_fct(layer_ind, t, x_t, (h_t, c_t), int_weights,
                                     curr_cm_weights, ckpt_id, is_last_step)

            if self.bptt_depth != -1:
                if t < (seq_length - self.bptt_depth):
                    # Detach hidden/output states, such that we don't backprop
                    # through these timesteps.
                    h_t = h_t.detach()
                    c_t = c_t.detach()

            # FIXME Solution is a bit ugly. For an LSTM, the hidden state is
            # also the output whereas a normal RNN has a separate output.
            if self._use_lstm:
                outputs.append(h_t)
                hiddens.append(c_t)
            else:
                outputs.append(c_t)
                hiddens.append(h_t)

        return torch.stack(outputs), torch.stack(hiddens)


    def basic_rnn_step(self, d, t, x_t, h_t, int_weights, cm_weights, ckpt_id,
                       is_last_step):
        """Perform vanilla rnn pass from inputs to hidden units.

        Apply context modulation if necessary (i.e. if ``cm_weights`` is
        not ``None``).

        This function implements a step of an
        `Elman RNN <https://en.wikipedia.org/wiki/\
Recurrent_neural_network#Elman_networks_and_Jordan_networks>`__.

        Note:
            We made the following design choice regarding context-modulation.
            In contrast to the LSTM, the Elman network layer consists of "two
            steps", updating the hidden state and computing an output based
            on this hidden state. To be fair, context-mod should influence both
            these "layers". Therefore, we apply context-mod twice, but using the
            same weights. This of course assumes that the hidden state and
            output vector have the same dimensionality.

        Args:
            d (int): Index of the layer.
            t (int): Current timestep.
            x_t: Tensor of size ``[batch_size, n_hidden_prev]`` with inputs.
            h_t (tuple): Tuple of length 2, containing two tensors of size
                ``[batch_size, n_hidden]`` with previous hidden states ``h`` and
                and previous outputs ``y``.

                Note:
                    The previous outputs ``y`` are ignored by this method, since
                    they are not required in an Elman RNN step.
            int_weights: See docstring of method :meth:`compute_hidden_states`.
            cm_weights (list): The weights of the context-mod layer, if context-
                mod should be applied.
            ckpt_id: See docstring of method :meth:`compute_hidden_states`.
            is_last_step (bool): Whether the current time step is the last one.

        Returns:
            (tuple): Tuple containing:

            - **h_t** (torch.Tensor): The tensor ``h_t`` of size
              ``[batch_size, n_hidden]`` with the new hidden state.
            - **y_t** (torch.Tensor): The tensor ``y_t`` of size
              ``[batch_size, n_hidden]`` with the new cell state.
        """
        h_t = h_t[0]

        use_cm = self._use_context_mod and d < self._num_rec_cm_layers
        # Determine the index of the hidden context mod layer.
        # Number of cm-layers per recurrent layer.
        n_cm_per_rec = self._context_mod_num_ts if \
            self._context_mod_num_ts != -1 and \
                self._context_mod_separate_layers_per_ts else 1
        cm_idx = self._cm_rnn_start_ind + d * n_cm_per_rec
        if self._context_mod_num_ts != -1 and \
                self._context_mod_separate_layers_per_ts:
            cm_idx += t
 
        if self.has_bias:
            assert len(int_weights) == 6
            weight_ih = int_weights[0]
            bias_ih = int_weights[1]
            weight_hh = int_weights[2]
            bias_hh = int_weights[3]
        else:
            assert len(int_weights) == 3
            weight_ih = int_weights[0]
            bias_ih = None
            weight_hh = int_weights[1]
            bias_hh = None

        ###########################
        ### Update hidden state ###
        ###########################
        h_t = x_t @ weight_ih.t() + h_t @ weight_hh.t()
        if self.has_bias:
            h_t += bias_ih + bias_hh

        # Context-dependent modulation (pre-activation).
        if use_cm and not self._context_mod_post_activation:
            # Only apply context mod if you are in the last time step, or if
            # you want to apply it in every single time step (i.e. if
            # self._context_mod_last_step is False).
            if not self._context_mod_last_step or is_last_step:
                h_t = self._context_mod_layers[cm_idx].forward(h_t,
                    weights=cm_weights, ckpt_id=ckpt_id)

        # Compute activation.
        if self._a_fun is not None:
            h_t = self._a_fun(h_t)

        # Context-dependent modulation (post-activation).
        if use_cm and self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                h_t = self._context_mod_layers[cm_idx].forward(h_t,
                    weights=cm_weights, ckpt_id=ckpt_id)

        ######################
        ### Compute output ###
        ######################
        y_t = self.compute_basic_rnn_output(h_t, int_weights, use_cm, 
            cm_weights, cm_idx, ckpt_id, is_last_step)

        return h_t, y_t

    def compute_basic_rnn_output(self, h_t, int_weights, use_cm, cm_weights,
                                 cm_idx, ckpt_id, is_last_step):
        """Compute the output of a vanilla RNN given the hidden state.

        Args:
            (...): See docstring of method :meth:`basic_rnn_step`.
            use_cm (boolean): Whether context modulation is being used.
            cm_idx (int): Index of the context-mod layer.

        Returns:
            (torch.tensor): The output.
        """
        if self.has_bias:
            weight_ho = int_weights[4]
            bias_ho = int_weights[5]
        else:
            weight_ho = int_weights[2]
            bias_ho = None

        y_t = h_t @ weight_ho.t()
        if self.has_bias:
            y_t += bias_ho

        # Context-dependent modulation (pre-activation).
        if use_cm and not self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                y_t = self._context_mod_layers[cm_idx].forward(y_t,
                    weights=cm_weights, ckpt_id=ckpt_id)

        # Compute activation.
        if self._a_fun is not None:
            y_t = self._a_fun(y_t)

        # Context-dependent modulation (post-activation).
        if use_cm and self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                y_t = self._context_mod_layers[cm_idx].forward(y_t,
                    weights=cm_weights, ckpt_id=ckpt_id)
        return y_t

    def lstm_rnn_step(self, d, t, x_t, h_t, int_weights, cm_weights, ckpt_id,
                      is_last_step):
        """ Perform an LSTM pass from inputs to hidden units.

        Apply masks to the temporal sequence for computing the loss.
        Obtained from:

            https://mlexplained.com/2019/02/15/building-an-lstm-from-scratch-\
in-pytorch-lstms-in-depth-part-1/

        and:

            https://d2l.ai/chapter_recurrent-neural-networks/lstm.html

        Args:
            d (int): Index of the layer.
            t (int): Current timestep.
            x_t: Tensor of size ``[batch_size, n_inputs]`` with inputs.
            h_t (tuple): Tuple of length 2, containing two tensors of size
                ``[batch_size, n_hidden]`` with previous hidden states ``h`` and
                ``c``.
            int_weights: See docstring of method :meth:`basic_rnn_step`.
            cm_weights: See docstring of method :meth:`basic_rnn_step`.
            ckpt_id: See docstring of method :meth:`basic_rnn_step`.
            is_last_step (bool): See docstring of method :meth:`basic_rnn_step`.

        Returns:
            (tuple): Tuple containing:

            - **h_t** (torch.Tensor): The tensor ``h_t`` of size
              ``[batch_size, n_hidden]`` with the new hidden state.
            - **c_t** (torch.Tensor): The tensor ``c_t`` of size
              ``[batch_size, n_hidden]`` with the new cell state.
        """
        use_cm = self._use_context_mod and d < self._num_rec_cm_layers
        # Determine the index of the hidden context mod layer.
        # Number of cm-layers per recurrent layer.
        n_cm_per_rec = self._context_mod_num_ts if \
            self._context_mod_num_ts != -1 and \
                self._context_mod_separate_layers_per_ts else 1
        cm_idx = self._cm_rnn_start_ind + d * n_cm_per_rec
        if self._context_mod_num_ts != -1 and \
                self._context_mod_separate_layers_per_ts:
            cm_idx += t

        c_t = h_t[1]
        h_t = h_t[0]
        HS = self._rnn_layers[d]

        if self._has_bias:
            assert len(int_weights) == 4
            weight_ih = int_weights[0]
            bias_ih = int_weights[1]
            weight_hh = int_weights[2]
            bias_hh = int_weights[3]
        else:
            assert len(int_weights) == 2
            weight_ih = int_weights[0]
            bias_ih = None
            weight_hh = int_weights[1]
            bias_hh = None

        # Compute total pre-activation input.
        gates = x_t @ weight_ih.t() + h_t @ weight_hh.t()
        if self.has_bias:
            gates += bias_ih + bias_hh

        i_t = gates[:, :HS]
        f_t = gates[:, HS:HS*2]
        g_t = gates[:, HS*2:HS*3]
        o_t = gates[:, HS*3:]

        # Compute activation.
        i_t = torch.sigmoid(i_t) # input
        f_t = torch.sigmoid(f_t) # forget
        g_t = torch.tanh(g_t)
        o_t = torch.sigmoid(o_t) # output

        # Compute c states.
        c_t = f_t * c_t + i_t * g_t

        # Note, we don't want to context-modulate the internal state c_t.
        # Otherwise, it might explode over timesteps since it wouldn't be
        # limited to [-1, 1] anymore. Instead, we only modulate the current
        # state (which is used to compute the current output h_t.
        c_t_mod = c_t

        # Context-dependent modulation (pre-activation).
        if use_cm and not self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                c_t_mod = self._context_mod_layers[cm_idx].forward(c_t_mod,
                    weights=cm_weights, ckpt_id=ckpt_id)

        # Compute h states.
        if self._a_fun is not None:
            h_t = o_t * self._a_fun(c_t_mod)
        else:
            h_t = o_t * c_t_mod

        # Context-dependent modulation (post-activation).
        # FIXME Christian: Shouldn't we apply the output gate `o_t` after
        # applying post-activation context-mod?
        if use_cm and self._context_mod_post_activation:
            if not self._context_mod_last_step or is_last_step:
                h_t = self._context_mod_layers[cm_idx].forward(h_t,
                    weights=cm_weights, ckpt_id=ckpt_id)

        return h_t, c_t

    def distillation_targets(self):
        """Targets to be distilled after training.

        See docstring of abstract super method
        :meth:`mnets.mnet_interface.MainNetInterface.distillation_targets`.

        This network does not have any distillation targets.

        Returns:
            ``None``
        """
        return None

    def get_output_weight_mask(self, out_inds=None, device=None):
        """Get masks to select output weights.

        See docstring of overwritten super method
        :meth:`mnets.mnet_interface.MainNetInterface.get_output_weight_mask`.
        """
        if len(self._fc_layers) > 0:
            return MainNetInterface.get_output_weight_mask(self,
                out_inds=out_inds, device=device)

        # TODO Output layer is recurrent. Hence, we have to properly handle
        # which weights contribute solely to certain output activations.
        raise NotImplementedError()

    def get_cm_weights(self):
        """Get internal maintained weights that are associated with context-
        modulation.

        Returns:
            (list): List of weights from
            :attr:`mnets.mnet_interface.MainNetInterface.internal_params` that
            are belonging to context-mod layers.
        """
        n_cm = self._num_context_mod_shapes()

        if n_cm == 0 or self._context_mod_no_weights:
            raise ValueError('Network maintains no context-modulation weights.')

        return self.internal_params[:n_cm]

    def get_non_cm_weights(self):
        """Get internal weights that are not associated with context-modulation.

        Returns:
            (list): List of weights from
            :attr:`mnets.mnet_interface.MainNetInterface.internal_params` that
            are not belonging to context-mod layers.
        """
        n_cm = 0 if self._context_mod_no_weights else \
            self._num_context_mod_shapes()

        return self.internal_params[n_cm:]

    def get_cm_inds(self):
        """Get the indices of
        :attr:`mnets.mnet_interface.MainNetInterface.param_shapes` that are
        associated with context-modulation.

        Returns:
            (list): List of integers representing indices of
            :attr:`mnets.mnet_interface.MainNetInterface.param_shapes`.
        """
        n_cm = self._num_context_mod_shapes()

        ret = []

        for i, meta in enumerate(self.param_shapes_meta):
            if meta['name'] == 'cm_shift' or meta['name'] == 'cm_scale':
                ret.append(i)
        assert n_cm == len(ret)

        return ret

    def init_hh_weights_orthogonal(self):
        """Initialize hidden-to-hidden weights orthogonally.

        This method will overwrite the hidden-to-hidden weights of recurrent
        layers.
        """
        for meta in self.param_shapes_meta:
            if meta['name'] == 'weight' and 'info' in meta.keys() and \
                    meta['info'] == 'hh' and meta['index'] != -1:
                print('Initializing hidden-to-hidden weights of recurrent ' +
                      'layer %d orthogonally.' % meta['layer'])
                W = self.internal_params[meta['index']]
                # LSTM weight matrices are stored such that the hidden-to-hidden 
                # matrices for the 4 gates are concatenated.
                if self.use_lstm:
                    out_dim, _ = W.shape
                    assert out_dim % 4 == 0
                    fs = out_dim // 4

                    W1 = W[:fs, :]
                    W2 = W[fs:2*fs, :]
                    W3 = W[2*fs:3*fs, :]
                    W4 = W[3*fs:, :]

                    torch.nn.init.orthogonal_(W1.data)
                    torch.nn.init.orthogonal_(W2.data)
                    torch.nn.init.orthogonal_(W3.data)
                    torch.nn.init.orthogonal_(W4.data)

                    # Sanity check to show that the init on partial matrices
                    # propagates back to the original tensor.
                    assert W[0,0] == W1[0,0]
                else:
                    torch.nn.init.orthogonal_(W.data)

    def _internal_weight_shapes(self):
        """Compute the tensor shapes of all internal weights (i.e., those not
        associated with context-modulation).

        Returns:
            (list): A list of list of integers, denoting the shapes of the
            individual parameter tensors.
        """
        coeff = 4 if self._use_lstm else 1
        shapes = []

        # Initial fully-connected layers.
        prev_dim = self._n_in
        for n_fc in self._fc_layers_pre:
            shapes.append([n_fc, prev_dim])
            if self._use_bias:
                shapes.append([n_fc])

            prev_dim = n_fc

        # Recurrent layers.
        for n_rec in self._rnn_layers:
            # Input-to-hidden
            shapes.append([n_rec*coeff, prev_dim])
            if self._use_bias:
                shapes.append([n_rec*coeff])

            # Hidden-to-hidden
            shapes.append([n_rec*coeff, n_rec])
            if self._use_bias:
                shapes.append([n_rec*coeff])

            if not self._use_lstm:
                # Hidden-to-output
                shapes.append([n_rec, n_rec])
                if self._use_bias:
                    shapes.append([n_rec])

            prev_dim = n_rec

        # Fully-connected layers.
        for n_fc in self._fc_layers:
            shapes.append([n_fc, prev_dim])
            if self._use_bias:
                shapes.append([n_fc])

            prev_dim = n_fc

        return shapes

if __name__ == '__main__':
    pass
#!/usr/bin/env python3 # Copyright 2019 <NAME> # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #    http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # title          :mnets/simple_rnn.py # author         :mc, be # contact        :mariacer, <EMAIL> # created        :10/28/2019 # version        :1.0 # python_version :3.6.8 SimpleRNN ---------  Implementation of a simple recurrent neural network that has stacked vanilla RNN or LSTM layers that are optionally enclosed by fully-connected layers.  An example usage is as a main model, where the main weights are initialized and protected by a method such as EWC, and the context-modulation patterns of the neurons are produced by an external hypernetwork. Implementation of a simple RNN.      This is a simple recurrent network, that receives input vector     :math:`\mathbf{x}` and outputs a vector :math:`\mathbf{y}` of real values.      Note:         The output is non-linear if the last layer is recurrent! Otherwise,         logits are returned (cmp. attribute         :attr:`mnets.mnet_interface.MainNetInterface.has_fc_out`).      Attributes:         bptt_depth (int): The truncation depth for backprop through time.             If ``-1``, backprop through time (BPTT) will unroll all timesteps             present in the input. Otherwise, the forward pass will detach the             RNN hidden states smaller or equal than             ``num_timesteps - bptt_depth`` timesteps, resulting in truncated             BPTT (T-BPTT).         num_rec_layers (int): Number of recurrent layers in this network (i.e.,             length of constructor argument ``rnn_layers``).         use_lstm (bool): See constructor argument ``use_lstm``.      Args:         n_in (int): Number of inputs.         rnn_layers (list or tuple): List of integers. Each entry denotes the             size of a recurrent layer. Recurrent layers will simply be stacked             as layers of this network.              If ``fc_layers_pre`` is empty, then the recurrent layers are the             initial layers.             If ``fc_layers`` is empty, then the last entry of this list will             denote the output size.              Note:                 This list may never be empty.         fc_layers_pre (list or tuple): List of integers. Before the recurrent             layers a set of fully-connected layers may be added. This might be             specially useful when constructing recurrent autoencoders. The             entries of this list will denote the sizes of those layers.              If ``fc_layers_pre`` is not empty, its first entry will denote the             input size of this network.         fc_layers (list or tuple): List of integers. After the recurrent layers,             a set of fully-connected layers is added. The entries of this list             will denote the sizes of those layers.              If ``fc_layers`` is not empty, its last entry will denote the output             size of this network.         activation: The nonlinearity used in hidden layers.         use_lstm (bool): If set to `True``, the recurrent layers will be LSTM             layers.         use_bias (bool): Whether layers may have bias terms.         no_weights (bool): If set to ``True``, no trainable parameters will be             constructed, i.e., weights are assumed to be produced ad-hoc             by a hypernetwork and passed to the :meth:`forward` method.         init_weights (list, optional): This option is for convinience reasons.             The option expects a list of parameter values that are used to             initialize the network weights. As such, it provides a             convinient way of initializing a network with a weight draw             produced by the hypernetwork.              Note, internal weights (see              :attr:`mnets.mnet_interface.MainNetInterface.weights`) will be             affected by this argument only.         kaiming_rnn_init (bool): By default, PyTorch initializes its recurrent             layers uniformly with an interval defined by the square-root of the             inverse of the layer size.              If this option is enabled, then the recurrent layers will be             initialized using the kaiming init as implemented by the function             :func:`utils.torch_utils.init_params`.         context_mod_last_step (bool): Whether context modulation is applied             at the last time step os a recurrent layer only. If ``False``,             context modulation is applied at every time step.              Note:                 This option only applies if ``use_context_mod`` is ``True``.         context_mod_num_ts (int, optional): The maximum number of timesteps.             If specified, context-modulation with a different set of weights is             applied at every timestep. If ``context_mod_separate_layers_per_ts``             is ``True``, then a separate context-mod layer per timestep will be             created. Otherwise, a single context-mod layer is created, but the             expected parameter shapes for this layer are             ``[context_mod_num_ts, *context_mod_shape]``.              Note:                 This option only applies if ``use_context_mod`` is ``True``.         context_mod_separate_layers_per_ts (bool): If specified, a separate             context-mod layer per timestep is created (required if             ``context_mod_no_weights`` is ``False``).                          Note:                 Only applies if ``context_mod_num_ts`` is specified.         verbose (bool): Whether to print information (e.g., the number of             weights) during the construction of the network.         **kwargs: Keyword arguments regarding context modulation. This class             can process the same context-modulation related arguments as class             :class:`mnets.mlp.MLP` (plus the additional ones noted above). # FIXME find a way using super to handle multiple inheritance. # FIXME Arbitrary restriction. #n_out = rnn_layers[-1] #n_out = fc_layers[-1] ### Parse or set context-mod arguments ### # Context-mod options specific to RNNs # FIXME We have to specify this option even if # `context_mod_separate_layers_per_ts` is False (in order to set # sensible parameter shapes). However, the forward method can deal with # an arbitrary timestep length. # More appropriate naming of option. ### Parse or set context-mod arguments - DONE ### # We need to make sure that the last 2 entries of `weights` correspond # to the weight matrix and bias vector of the last layer. # Note, recurrent layers always use non-linearities and their activities # are squashed by a non-linearity (otherwise, internal states could # vanish/explode with increasing sequence length). ################################################# ### Define and initialize context mod weights ### ################################################# # The context-mod layers consist of sequential layers ordered as: # - initial fully-connected layers if len(fc_layers_pre)>0 # - recurrent layers # - final fully-connected layers if len(fc_layers)>0 # Gather sizes of all activation vectors within the network that # will be subject to context-modulation. # We reserve layer zero for input context-mod. Otherwise, there # is no layer zero. # We use odd numbers for actual layers and even number for all # context-mod layers. # Add initial fully-connected context-mod layers. # Add recurrent context-mod layers. # Only a single context-mod layer will be added, but we # directly edit the correponding `param_shape` later. # Add final fully-connected context-mod layers. # In this case, there is only one context-mod layer for each # recurrent layer, but we want to have separate weights per # timestep. # Hence, we adapt the expected parameter shape, such that we # get a different set of weights per timestep. This will be # split into multiple weights that are succesively fed into the # same layer inside the forward method. ######################## ### Internal weights ### ######################## Define the weights and shapes of the fully-connected layers.              Args:                 fc_layers (list): The list of fully-connected layer dimensions.                 prev_dim (int): The output size of the previous layer.                 num_prev_layers (int): The number of upstream layers to the                      current one (a layer with its corresponding                     context-mod layer(s) count as one layer). Count should                     start at ``1``.              Returns:                 (int): The output size of the last fully-connected layer                 considered here. # FIXME We should instead build an MLP instance. But then we still # have to adapt all attributes accordingly. # Odd numbers ### Initial fully-connected layers. ### Recurrent layers. # Input-to-hidden # Hidden-to-hidden # Hidden-to-output. # Note, for an LSTM cell, the hidden state vector is also the # output vector. # Odd numbers ### Fully-connected layers. ### Initialize weights. # Note, Pytorch applies a uniform init to its recurrent layers, as # defined here: # https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py#L155 # Recurrent layer weights. # FC layer weights. Getter for attribute :attr:`bptt_depth`. Setter for attribute :attr:`bptt_depth`. Getter for read-only attribute :attr:`num_rec_layers`. Getter for read-only attribute :attr:`use_lstm`. Split context-mod weights per context-mod layer.          Args:             cm_weights (torch.Tensor): All context modulation weights.             condition (optional, int): If provided, then this argument will be                 passed as argument ``ckpt_id`` to the method                 :meth:`utils.context_mod_layer.ContextModLayer.forward`.             num_ts (int): The length of the sequences.          Returns:             (Tuple): Where the tuple contains:              - **cm_inputs_weights**: The cm input weights.             - **cm_fc_pre_layer_weights**: The cm pre-recurrent weights.             - **cm_rec_layer_weights**: The cm recurrent weights.             - **cm_fc_layer_weights**: The cm post-recurrent weights.             - **n_cm_rec**: The number of recurrent cm layers.             - **cmod_cond**: The context-mod condition. # Note, an mnet layer might contain multiple context-mod layers # (a recurrent layer can have a separate context-mod layer per # timestep). # Number of cm-layers per recurrent layer. # Index of recurrent layer. # We need to split the context-mod weights in the following case, # as they are currently just stacked on top of each other. # The first dimension are the weights of this layer per # timestep. # Now we gather all these chunks to assemble the weights # needed per timestep (as if # `_context_mod_separate_layers_per_t` were True). # Note, the last layer does not necessarily have context-mod # (depending on `self._context_mod_outputs`). ####################### ### Parse condition ### ####################### # Note, the cm layer will ignore the cmod condition if weights # are passed. # FIXME Find a more elegant solution. Split internal weights per layer.          Args:             int_weights (torch.Tensor): All internal weights.          Returns:             (Tuple): Where the tuple contains:              - **fc_pre_w_weights**: The pre-recurrent w weights.             - **fc_pre_b_weights**: The pre-recurrent b weights.             - **rec_weights**: The recurrent weights.             - **fc_w_weights**:The post-recurrent w weights.             - **fc_b_weights**: The post-recurrent b weights. # Number of pre-fc weights in total. # Number of weights per recurrent layer. # fc pre weights # recurrent w # fc weights Split weights into internal and context-mod weights.          Extract which weights should be used,  I.e., are we using internally         maintained weights or externally given ones or are we even mixing         between these groups.          Args:             weights (torch.Tensor): All weights.          Returns:             (Tuple): Where the tuple contains:              - **int_weights**: The internal weights.             - **cm_weights**: The context-mod weights. ### FIXME Code copied from MLP its `forward` method ### # Make sure cm_weights are either `None` or have the correct dimensions. # Note, context-mod weights might have different shapes, as they # may be parametrized on a per-sample basis. ### FIXME Code copied until here ### Compute the output :math:`y` of this network given the input         :math:`x`.          Args:             (....): See docstring of method                 :meth:`mnets.mnet_interface.MainNetInterface.forward`. We                 provide some more specific information below.             weights (list or dict): See argument ``weights`` of method                 :meth:`mnets.mlp.MLP.forward`.             condition (optional, int): If provided, then this argument will be                 passed as argument ``ckpt_id`` to the method                 :meth:`utils.context_mod_layer.ContextModLayer.forward`.             return_hidden (bool, optional): If ``True``, all hidden activations                 of fully-connected and recurrent layers (where we defined                 :math:`y_t` as hidden state of vannila RNN layers as these are                 the layer outputs passed to the next layer) are returned.                 recurrent hidden activations will be returned.             return_hidden_int (bool, optional): If ``True``, in addition to                 ``hidden``, an additional variable ``hidden_int`` is returned                 containing the internal hidden states of recurrent layers (i.e.,                 the cell states :math:`c_t` for LSTMs and the actual hidden                 state :math:`h_t` for Elman layers) are returned.          Returns:             (torch.Tensor or tuple): Where the tuple is containing:              - **output** (torch.Tensor): The output of the network.             - **hidden** (list): If ``return_hidden`` is ``True``, then the               hidden activities of each layer are returned, which have the shape               ``(seq_length, batch_size, n_hidden)``.             - **hidden_int**: If ``return_hidden_int`` is ``True``, then in               addition to ``hidden`` a tensor ``hidden_int`` is returned               containing internal hidden states of recurrent layers. ####################### ### Extract weights ### ####################### # Extract which weights should be used. ### Split context-mod weights per context-mod layer. ### Extract internal weights. ########################### ### Forward Computation ### ########################### # Apply context modulation in the inputs. ### Initial fully-connected layer activities. ### Recurrent layer activities. # the internal hidden activations ### Fully-connected layer activities. # FIXME quite ugly # The last element is the output activity. Compute the forward pass through the fully-connected layers.          This method also appends activations to ``ret_hidden``.          Args:             h (torch.Tensor): The input from the previous layer.             fc_w_weights (list): The weights for the fc layers.             fc_b_weights (list): The biases for the fc layers.             num_fc_cm_layers (int): The number of context-modulation                 layers associated with this set of fully-connected layers.             cm_fc_layer_weights (list): The context-modulation weights                 associated with the current layers.             cm_offset (int): The index to access the correct context-mod                 layers.             cmod_cond (bool): Some condition to perform context modulation.             is_post_fc (bool); Whether those layers are applied as last                 layers of the network. In this case, there will be no                 activation applied to the last layer outputs.             ret_hidden (list or None): The hidden recurrent activations.          Return:             (Tuple): Tuple containing:              - **ret_hidden**: The hidden recurrent activations.             - **h**: Transformed activation ``h``. # Compute output. # Context-dependent modulation (pre-activation). # Non-linearity # Note, non-linearity is not applied to outputs of the network. # Context-dependent modulation (post-activation). Compute the hidden states for the recurrent layer ``layer_ind`` from         a sequence of inputs :math:`x`.          If so specified, context modulation is applied before or after the         nonlinearities.          Args:             x: The inputs :math:`x` to the layer. :math:`x` has shape                 ``[sequence_len, batch_size, n_hidden_prev]``.             layer_ind (int): Index of the layer.             int_weights: Internal weights associated with this recurrent layer.             cm_weights: Context modulation weights.             ckpt_id: Will be passed as option ``ckpt_id`` to method                 :meth:`utils.context_mod_layer.ContextModLayer.forward` if                 context-mod layers are used.             h_0 (torch.Tensor, optional): The initial state for :math:`h`.             c_0 (torch.Tensor, optional): The initial state for :math:`c`. Note                 that for LSTMs, if the initial state is to be defined, this                 variable is necessary also, not only :math:`h_0`, whereas for                 vanilla RNNs it is enough to provide :math:`h_0` as :math:`c_0`                 represents the output of the layer and it can be easily computed                 from `h_0`.          Returns:             (tuple): Tuple containing:              - **outputs** (torch.Tensor): The sequence of visible hidden states               given the input. It has shape                ``[sequence_len, batch_size, n_hidden]``.             - **hiddens** (torch.Tensor): The sequence of hidden states given               the input. For LSTMs, this corresponds to :math:`c`.               It has shape ``[sequence_len, batch_size, n_hidden]``. # Generate initial hidden states. # Note that for LSTMs h_0 is the hidden state and output vector whereas # c_0 is the internal cell state vector. # For a vanilla RNN h_0 is the hidden state whereas c_0 is the output # vector. # If we want to apply context modulation in each time step, we need # to split the input sequence and call pytorch function at every # time step. # Compute the actual rnn step (either vanilla or LSTM, depending on # the flag self._use_lstm). # Detach hidden/output states, such that we don't backprop # through these timesteps. # FIXME Solution is a bit ugly. For an LSTM, the hidden state is # also the output whereas a normal RNN has a separate output. Perform vanilla rnn pass from inputs to hidden units.          Apply context modulation if necessary (i.e. if ``cm_weights`` is         not ``None``).          This function implements a step of an         `Elman RNN <https://en.wikipedia.org/wiki/\ Recurrent_neural_network#Elman_networks_and_Jordan_networks>`__.          Note:             We made the following design choice regarding context-modulation.             In contrast to the LSTM, the Elman network layer consists of "two             steps", updating the hidden state and computing an output based             on this hidden state. To be fair, context-mod should influence both             these "layers". Therefore, we apply context-mod twice, but using the             same weights. This of course assumes that the hidden state and             output vector have the same dimensionality.          Args:             d (int): Index of the layer.             t (int): Current timestep.             x_t: Tensor of size ``[batch_size, n_hidden_prev]`` with inputs.             h_t (tuple): Tuple of length 2, containing two tensors of size                 ``[batch_size, n_hidden]`` with previous hidden states ``h`` and                 and previous outputs ``y``.                  Note:                     The previous outputs ``y`` are ignored by this method, since                     they are not required in an Elman RNN step.             int_weights: See docstring of method :meth:`compute_hidden_states`.             cm_weights (list): The weights of the context-mod layer, if context-                 mod should be applied.             ckpt_id: See docstring of method :meth:`compute_hidden_states`.             is_last_step (bool): Whether the current time step is the last one.          Returns:             (tuple): Tuple containing:              - **h_t** (torch.Tensor): The tensor ``h_t`` of size               ``[batch_size, n_hidden]`` with the new hidden state.             - **y_t** (torch.Tensor): The tensor ``y_t`` of size               ``[batch_size, n_hidden]`` with the new cell state. # Determine the index of the hidden context mod layer. # Number of cm-layers per recurrent layer. ########################### ### Update hidden state ### ########################### # Context-dependent modulation (pre-activation). # Only apply context mod if you are in the last time step, or if # you want to apply it in every single time step (i.e. if # self._context_mod_last_step is False). # Compute activation. # Context-dependent modulation (post-activation). ###################### ### Compute output ### ###################### Compute the output of a vanilla RNN given the hidden state.          Args:             (...): See docstring of method :meth:`basic_rnn_step`.             use_cm (boolean): Whether context modulation is being used.             cm_idx (int): Index of the context-mod layer.          Returns:             (torch.tensor): The output. # Context-dependent modulation (pre-activation). # Compute activation. # Context-dependent modulation (post-activation). Perform an LSTM pass from inputs to hidden units.          Apply masks to the temporal sequence for computing the loss.         Obtained from:              https://mlexplained.com/2019/02/15/building-an-lstm-from-scratch-\ in-pytorch-lstms-in-depth-part-1/          and:              https://d2l.ai/chapter_recurrent-neural-networks/lstm.html          Args:             d (int): Index of the layer.             t (int): Current timestep.             x_t: Tensor of size ``[batch_size, n_inputs]`` with inputs.             h_t (tuple): Tuple of length 2, containing two tensors of size                 ``[batch_size, n_hidden]`` with previous hidden states ``h`` and                 ``c``.             int_weights: See docstring of method :meth:`basic_rnn_step`.             cm_weights: See docstring of method :meth:`basic_rnn_step`.             ckpt_id: See docstring of method :meth:`basic_rnn_step`.             is_last_step (bool): See docstring of method :meth:`basic_rnn_step`.          Returns:             (tuple): Tuple containing:              - **h_t** (torch.Tensor): The tensor ``h_t`` of size               ``[batch_size, n_hidden]`` with the new hidden state.             - **c_t** (torch.Tensor): The tensor ``c_t`` of size               ``[batch_size, n_hidden]`` with the new cell state. # Determine the index of the hidden context mod layer. # Number of cm-layers per recurrent layer. # Compute total pre-activation input. # Compute activation. # input # forget # output # Compute c states. # Note, we don't want to context-modulate the internal state c_t. # Otherwise, it might explode over timesteps since it wouldn't be # limited to [-1, 1] anymore. Instead, we only modulate the current # state (which is used to compute the current output h_t. # Context-dependent modulation (pre-activation). # Compute h states. # Context-dependent modulation (post-activation). # FIXME Christian: Shouldn't we apply the output gate `o_t` after # applying post-activation context-mod? Targets to be distilled after training.          See docstring of abstract super method         :meth:`mnets.mnet_interface.MainNetInterface.distillation_targets`.          This network does not have any distillation targets.          Returns:             ``None`` Get masks to select output weights.          See docstring of overwritten super method         :meth:`mnets.mnet_interface.MainNetInterface.get_output_weight_mask`. # TODO Output layer is recurrent. Hence, we have to properly handle # which weights contribute solely to certain output activations. Get internal maintained weights that are associated with context-         modulation.          Returns:             (list): List of weights from             :attr:`mnets.mnet_interface.MainNetInterface.internal_params` that             are belonging to context-mod layers. Get internal weights that are not associated with context-modulation.          Returns:             (list): List of weights from             :attr:`mnets.mnet_interface.MainNetInterface.internal_params` that             are not belonging to context-mod layers. Get the indices of         :attr:`mnets.mnet_interface.MainNetInterface.param_shapes` that are         associated with context-modulation.          Returns:             (list): List of integers representing indices of             :attr:`mnets.mnet_interface.MainNetInterface.param_shapes`. Initialize hidden-to-hidden weights orthogonally.          This method will overwrite the hidden-to-hidden weights of recurrent         layers. # LSTM weight matrices are stored such that the hidden-to-hidden # matrices for the 4 gates are concatenated. # Sanity check to show that the init on partial matrices # propagates back to the original tensor. Compute the tensor shapes of all internal weights (i.e., those not         associated with context-modulation).          Returns:             (list): A list of list of integers, denoting the shapes of the             individual parameter tensors. # Initial fully-connected layers. # Recurrent layers. # Input-to-hidden # Hidden-to-hidden # Hidden-to-output # Fully-connected layers.