popgym.baselines.models.linear_attention
========================================

.. py:module:: popgym.baselines.models.linear_attention


Classes
-------

.. autoapisummary::

   popgym.baselines.models.linear_attention.Phi
   popgym.baselines.models.linear_attention.LinearAttentionBlock


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

.. py:class:: Phi

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


   .. py:method:: forward(x)


.. py:class:: LinearAttentionBlock(input_size: int, hidden_size: int, S_aggregator: str = 'sum', Z_aggregator: str = 'sum', feed_forward=True, residual=True)

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


   The building block from the Linear Transformers are Secretly RNNs Paper. This is
   a form of linear transformer.

   Inputs:
       input_size: Size of input feature dim
       hidden_size: Size of key/query/value space
       S_aggregator: Which type of aggregation to use for the numerator (S term)
       Z_aggregator: Which type of aggregation to use for the denominator (Z term)
       feed_forward: Whether to apply a perceptron to the output
       residual: Whether to apply a residual connection from input to output


   .. py:attribute:: key


   .. py:attribute:: query


   .. py:attribute:: value


   .. py:attribute:: norm


   .. py:attribute:: phi


   .. py:attribute:: S_aggregator


   .. py:attribute:: Z_aggregator


   .. py:attribute:: feed_forward
      :value: True



   .. py:attribute:: residual
      :value: True



   .. py:method:: forward(x: torch.Tensor, state: List[torch.Tensor]) -> Tuple[torch.Tensor, List[torch.Tensor]]

      Input:
          x: [B, T, F]
          state: Tuple[
              [B, 1, D, D],
              [B, 1, D]
          ]
      Output:
          y: [B, T, D]
          state: Tuple[
              [B, 1, D, D],
              [B, 1, D]
          ]