foscat.healpix_vit_skip
=======================

.. py:module:: foscat.healpix_vit_skip


Attributes
----------

.. autoapisummary::

   foscat.healpix_vit_skip.in_nside


Classes
-------

.. autoapisummary::

   foscat.healpix_vit_skip.MLP
   foscat.healpix_vit_skip.HealpixViTSkip


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

.. py:class:: MLP(d: int, hidden_mult: int = 4, drop: float = 0.0)

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


   Base class for all neural network modules.

   Your models should also subclass this class.

   Modules can also contain other Modules, allowing them to be nested in
   a tree structure. You can assign the submodules as regular attributes::

       import torch.nn as nn
       import torch.nn.functional as F


       class Model(nn.Module):
           def __init__(self) -> None:
               super().__init__()
               self.conv1 = nn.Conv2d(1, 20, 5)
               self.conv2 = nn.Conv2d(20, 20, 5)

           def forward(self, x):
               x = F.relu(self.conv1(x))
               return F.relu(self.conv2(x))

   Submodules assigned in this way will be registered, and will also have their
   parameters converted when you call :meth:`to`, etc.

   .. note::
       As per the example above, an ``__init__()`` call to the parent class
       must be made before assignment on the child.

   :ivar training: Boolean represents whether this module is in training or
                   evaluation mode.
   :vartype training: bool


   .. py:attribute:: net


   .. py:method:: forward(x)


.. py:class:: HealpixViTSkip(*, in_nside: int, n_chan_in: int, level_dims: List[int], depth_token: int, num_heads_token: int, cell_ids: numpy.ndarray, task: Literal['regression', 'segmentation', 'global'] = 'regression', out_channels: int = 1, mlp_ratio_token: float = 4.0, KERNELSZ: int = 3, gauge_type: Literal['cosmo', 'phi'] = 'cosmo', G: int = 1, prefer_foscat_gpu: bool = True, dropout: float = 0.1, dtype: Literal['float32', 'float64'] = 'float32', pos_embed_per_level: bool = True)

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


   Base class for all neural network modules.

   Your models should also subclass this class.

   Modules can also contain other Modules, allowing them to be nested in
   a tree structure. You can assign the submodules as regular attributes::

       import torch.nn as nn
       import torch.nn.functional as F


       class Model(nn.Module):
           def __init__(self) -> None:
               super().__init__()
               self.conv1 = nn.Conv2d(1, 20, 5)
               self.conv2 = nn.Conv2d(20, 20, 5)

           def forward(self, x):
               x = F.relu(self.conv1(x))
               return F.relu(self.conv2(x))

   Submodules assigned in this way will be registered, and will also have their
   parameters converted when you call :meth:`to`, etc.

   .. note::
       As per the example above, an ``__init__()`` call to the parent class
       must be made before assignment on the child.

   :ivar training: Boolean represents whether this module is in training or
                   evaluation mode.
   :vartype training: bool


   .. py:attribute:: in_nside


   .. py:attribute:: n_chan_in


   .. py:attribute:: level_dims


   .. py:attribute:: token_down


   .. py:attribute:: C_fine


   .. py:attribute:: embed_dim


   .. py:attribute:: depth_token


   .. py:attribute:: num_heads_token


   .. py:attribute:: mlp_ratio_token


   .. py:attribute:: task
      :value: 'regression'



   .. py:attribute:: out_channels
      :value: 1



   .. py:attribute:: KERNELSZ
      :value: 3



   .. py:attribute:: gauge_type
      :value: 'cosmo'



   .. py:attribute:: G
      :value: 1



   .. py:attribute:: prefer_foscat_gpu
      :value: True



   .. py:attribute:: dropout


   .. py:attribute:: dtype
      :value: 'float32'



   .. py:attribute:: pos_embed_per_level
      :value: True



   .. py:attribute:: cell_ids_fine


   .. py:attribute:: f


   .. py:attribute:: hconv_levels
      :type:  List[foscat.SphericalStencil.SphericalStencil]
      :value: []



   .. py:attribute:: level_cell_ids
      :type:  List[numpy.ndarray]


   .. py:attribute:: token_nside


   .. py:attribute:: token_cell_ids


   .. py:attribute:: hconv_token


   .. py:attribute:: hconv_head


   .. py:attribute:: nsides_levels


   .. py:attribute:: ntokens_levels


   .. py:attribute:: pe_w1


   .. py:attribute:: pe_w2


   .. py:attribute:: pe_bn1


   .. py:attribute:: pe_bn2


   .. py:attribute:: enc_w1


   .. py:attribute:: enc_w2


   .. py:attribute:: enc_bn1


   .. py:attribute:: enc_bn2


   .. py:attribute:: temporal_encoders


   .. py:attribute:: n_tokens


   .. py:attribute:: pos_token


   .. py:attribute:: encoder_token


   .. py:attribute:: dec_q


   .. py:attribute:: dec_k


   .. py:attribute:: dec_v


   .. py:attribute:: dec_attn


   .. py:attribute:: dec_mlp


   .. py:attribute:: level_pos


   .. py:attribute:: dec_refine_w1


   .. py:attribute:: dec_refine_w2


   .. py:attribute:: dec_refine_bn1


   .. py:attribute:: dec_refine_bn2


   .. py:attribute:: runtime_device


   .. py:method:: forward(x: torch.Tensor, runtime_ids: Optional[numpy.ndarray] = None) -> torch.Tensor


.. py:data:: in_nside
   :value: 4