popgym.baselines.models.lmu
===========================

.. py:module:: popgym.baselines.models.lmu


Classes
-------

.. autoapisummary::

   popgym.baselines.models.lmu.LMUCell
   popgym.baselines.models.lmu.LMU
   popgym.baselines.models.lmu.LMUFFT


Functions
---------

.. autoapisummary::

   popgym.baselines.models.lmu.leCunUniform


Module Contents
---------------

.. py:function:: leCunUniform(tensor)

   LeCun Uniform Initializer


.. py:class:: LMUCell(input_size, hidden_size, memory_size, theta, learn_a=False, learn_b=False)

   Bases: :py:obj:`torch.nn.Module`


   LMU Cell.

   Parameters:
       input_size (int) :
           Size of the input vector (x_t)
       hidden_size (int) :
           Size of the hidden vector (h_t)
       memory_size (int) :
           Size of the memory vector (m_t)
       theta (int) :
           The number of timesteps in the sliding window that is
           represented using the LTI system
       learn_a (boolean) :
           Whether to learn the matrix A (default = False)
       learn_b (boolean) :
           Whether to learn the matrix B (default = False)


   .. py:attribute:: hidden_size


   .. py:attribute:: memory_size


   .. py:attribute:: f


   .. py:attribute:: e_x


   .. py:attribute:: e_h


   .. py:attribute:: e_m


   .. py:attribute:: W_x


   .. py:attribute:: W_h


   .. py:attribute:: W_m


   .. py:method:: initParameters()

      Initialize the cell's parameters.



   .. py:method:: stateSpaceMatrices(memory_size, theta)

      Returns the discretized state space matrices A and B.



   .. py:method:: forward(x, state)

      Parameters:
          x (torch.tensor):
              Input of size [batch_size, input_size]
          state (tuple):
              h (torch.tensor) : [batch_size, hidden_size]
              m (torch.tensor) : [batch_size, memory_size]



.. py:class:: LMU(input_size, hidden_size, memory_size, theta, learn_a=False, learn_b=False)

   Bases: :py:obj:`torch.nn.Module`


   LMU layer.

   Parameters:
       input_size (int) :
           Size of the input vector (x_t)
       hidden_size (int) :
           Size of the hidden vector (h_t)
       memory_size (int) :
           Size of the memory vector (m_t)
       theta (int) :
           The number of timesteps in the sliding window
           that is represented using the LTI system
       learn_a (boolean) :
           Whether to learn the matrix A (default = False)
       learn_b (boolean) :
           Whether to learn the matrix B (default = False)


   .. py:attribute:: hidden_size


   .. py:attribute:: memory_size


   .. py:attribute:: cell


   .. py:method:: forward(x, state=None)

      Parameters:
          x (torch.tensor):
              Input of size [batch_size, seq_len, input_size]
          state (tuple) : (default = None)
              h (torch.tensor) : [batch_size, hidden_size]
              m (torch.tensor) : [batch_size, memory_size]



.. py:class:: LMUFFT(input_size, hidden_size, memory_size, seq_len, theta)

   Bases: :py:obj:`torch.nn.Module`


   Parallelized LMU Layer.

   Parameters:
       input_size (int) :
           Size of the input vector (x_t)
       hidden_size (int) :
           Size of the hidden vector (h_t)
       memory_size (int) :
           Size of the memory vector (m_t)
       seq_len (int) :
           Size of the sequence length (n)
       theta (int) :
           The number of timesteps in the sliding
           window that is represented using the LTI system


   .. py:attribute:: hidden_size


   .. py:attribute:: memory_size


   .. py:attribute:: seq_len


   .. py:attribute:: theta


   .. py:attribute:: W_u


   .. py:attribute:: f_u


   .. py:attribute:: W_h


   .. py:attribute:: f_h


   .. py:method:: stateSpaceMatrices()

      Returns the discretized state space matrices A and B.



   .. py:method:: impulse()

      Returns the matrices H and the 1D Fourier transform of H (Equations
      23, 26 of the paper)



   .. py:method:: forward(x)

      Parameters:
          x (torch.tensor):
              Input of size [batch_size, seq_len, input_size]