emmi.modules.blocks.transformer_block
=====================================

.. py:module:: emmi.modules.blocks.transformer_block


Classes
-------

.. autoapisummary::

   emmi.modules.blocks.transformer_block.TransformerBlock


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

.. py:class:: TransformerBlock(config)

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


   A transformer block with a single attention layer and a feedforward layer.

   Initializes a transformer block.

   Args:



   .. py:attribute:: norm1


   .. py:attribute:: attention_block


   .. py:attribute:: ls1


   .. py:attribute:: drop_path1


   .. py:attribute:: norm2


   .. py:attribute:: mlp


   .. py:attribute:: ls2


   .. py:attribute:: drop_path2


   .. py:method:: forward(x, condition = None, attn_kwargs = None)

      Forward pass of the transformer block.

      :param x: Input tensor with shape (batch_size, seqlen/num_tokens, hidden_dim).
      :param condition: Conditioning vector. If provided, the attention and MLP will be scaled, shifted and gated
                        feature-wise with predicted values from this vector.
      :param attn_kwargs: Dict with arguments for the attention (such as the attention mask). Defaults to None.

      :returns: Tensor after the forward pass of the transformer block.