class_docstr: Encoder for the Masked Autoencoder model [0].
  
  Encodes patch embeddings. Code inspired by [1].
  
  - [0]: Masked Autoencoder, 2021, https://arxiv.org/abs/2111.06377
  - [1]: https://github.com/facebookresearch/mae
  
  Attributes:
      seq_length:
          Token sequence length, including the class token.
      num_layers:
          Number of transformer blocks.
      num_heads:
          Number of attention heads.
      hidden_dim:
          Dimension of the input and output tokens.
      mlp_dim:
          Dimension of the MLP in the transformer block.
      dropout:
          Percentage of elements set to zero after the MLP in the transformer.
      attention_dropout:
          Percentage of elements set to zero after the attention head.
docstr: Encode input tokens.
  
  Args:
      input:
          Batch of token sequences.
      idx_keep:
          Tensor with shape (batch_size, num_tokens_to_keep) where each
          entry is an index of the token to keep in the respective batch.
          If specified, only the indexed tokens will be encoded.
  
  Returns:
      Batch of encoded output tokens.
docstr_tok: ["Encode","input","tokens",".","Args",":","input",":","Batch","of","token","sequences",".","idx_keep",":","Tensor","with","shape","(","batch_size",",","num_tokens_to_keep",")","where","each","entry","is","an","index","of","the","token","to","keep","in","the","respective","batch",".","If","specified",",","only","the","indexed","tokens","will","be","encoded",".","Returns",":","Batch","of","encoded","output","tokens","."]
code: def forward(
      self, input: torch.Tensor, idx_keep: Optional[torch.Tensor] = None
  ) -> torch.Tensor:
      """Encode input tokens.
  
      Args:
          input:
              Batch of token sequences.
          idx_keep:
              Tensor with shape (batch_size, num_tokens_to_keep) where each
              entry is an index of the token to keep in the respective batch.
              If specified, only the indexed tokens will be encoded.
  
      Returns:
          Batch of encoded output tokens.
      """
      input = input + self.interpolate_pos_encoding(input)
      if idx_keep is not None:
          input = utils.get_at_index(input, idx_keep)
      return self.ln(self.layers(self.dropout(input)))
  
imports: from __future__ import annotations
  import math
  from functools import partial
  from typing import Callable, List, Optional
  import torch
  import torch.nn
  from torchvision.models import vision_transformer
  from torchvision.models.vision_transformer import ConvStemConfig
  from lightly.models import utils
label_desc: Use image node_id 3 for calling the MAEEncoder obj's underlying member method code with example usage: obj.forward(input, idx_keep) and returns: self