emmi_inference.models.modules.rope_frequency
============================================

.. py:module:: emmi_inference.models.modules.rope_frequency


Classes
-------

.. autoapisummary::

   emmi_inference.models.modules.rope_frequency.RopeFrequency


Functions
---------

.. autoapisummary::

   emmi_inference.models.modules.rope_frequency.maybe_autocast


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

.. py:function:: maybe_autocast(device, enabled, dtype = None)

.. py:class:: RopeFrequency(dim, ndim, max_wavelength = 10000.0, assert_positive = True)

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


   Creates frequencies for rotary embeddings (RoPE) from https://arxiv.org/abs/2104.09864 for variable positions.

   :param dim: Dimensionality of frequencies (in transformers this should be the head dimension).
   :param ndim: Dimensionality of the coordinates (e.g., 2 for 2D coordinates, 3 for 3D coordinates).
   :param max_wavelength: Theta parameter for the transformer sine/cosine embedding. Default: 10000.0
   :param assert_positive: Makes sure that coordinates were rescaled to be positive only. Default: True

   Initialize internal Module state, shared by both nn.Module and ScriptModule.


   .. py:attribute:: dim


   .. py:attribute:: ndim


   .. py:attribute:: ndim_padding


   .. py:attribute:: sincos_padding


   .. py:attribute:: max_wavelength
      :value: 10000.0



   .. py:attribute:: padding


   .. py:attribute:: assert_positive
      :value: True



   .. py:method:: forward(coords)