code: def forward(
      self,
      batched_images: Union[
          List[Tensor], List[List[Tensor]]
      ],  # assume different resolution images already grouped correctly
      group_images=False,
      group_max_seq_len=2048,
  ):
      p, c, device, has_token_dropout = (
          self.patch_size,
          self.channels,
          self.device,
          exists(self.calc_token_dropout),
      )
  
      arange = partial(torch.arange, device=device)
      pad_sequence = partial(orig_pad_sequence, batch_first=True)
  
      # auto pack if specified
  
      if group_images:
          batched_images = group_images_by_max_seq_len(
              batched_images,
              patch_size=self.patch_size,
              calc_token_dropout=self.calc_token_dropout,
              max_seq_len=group_max_seq_len,
          )
  
      # process images into variable lengthed sequences with attention mask
  
      num_images = []
      batched_sequences = []
      batched_positions = []
      batched_image_ids = []
  
      for images in batched_images:
          num_images.append(len(images))
  
          sequences = []
          positions = []
          image_ids = torch.empty((0,), device=device, dtype=torch.long)
  
          for image_id, image in enumerate(images):
              assert image.ndim == 3 and image.shape[0] == c
              image_dims = image.shape[-2:]
              assert all(
                  [divisible_by(dim, p) for dim in image_dims]
              ), f"height and width {image_dims} of images must be divisible by patch size {p}"
  
              ph, pw = map(lambda dim: dim // p, image_dims)
  
              pos = torch.stack(
                  torch.meshgrid(
                      (arange(ph), arange(pw)), indexing="ij"
                  ),
                  dim=-1,
              )
  
              pos = rearrange(pos, "h w c -> (h w) c")
              seq = rearrange(
                  image,
                  "c (h p1) (w p2) -> (h w) (c p1 p2)",
                  p1=p,
                  p2=p,
              )
  
              seq_len = seq.shape[-2]
  
              if has_token_dropout:
                  token_dropout = self.calc_token_dropout(*image_dims)
                  num_keep = max(1, int(seq_len * (1 - token_dropout)))
                  keep_indices = (
                      torch.randn((seq_len,), device=device)
                      .topk(num_keep, dim=-1)
                      .indices
                  )
  
                  seq = seq[keep_indices]
                  pos = pos[keep_indices]
  
              image_ids = F.pad(
                  image_ids, (0, seq.shape[-2]), value=image_id
              )
              sequences.append(seq)
              positions.append(pos)
  
          batched_image_ids.append(image_ids)
          batched_sequences.append(torch.cat(sequences, dim=0))
          batched_positions.append(torch.cat(positions, dim=0))
  
      # derive key padding mask
  
      lengths = torch.tensor(
          [seq.shape[-2] for seq in batched_sequences],
          device=device,
          dtype=torch.long,
      )
      max_length = arange(lengths.amax().item())
      key_pad_mask = rearrange(lengths, "b -> b 1") <= rearrange(
          max_length, "n -> 1 n"
      )
  
      # derive attention mask, and combine with key padding mask from above
  
      batched_image_ids = pad_sequence(batched_image_ids)
      attn_mask = rearrange(
          batched_image_ids, "b i -> b 1 i 1"
      ) == rearrange(batched_image_ids, "b j -> b 1 1 j")
      attn_mask = attn_mask & rearrange(key_pad_mask, "b j -> b 1 1 j")
  
      # combine patched images as well as the patched width / height positions for 2d positional embedding
  
      patches = pad_sequence(batched_sequences)
      patch_positions = pad_sequence(batched_positions)
  
      # need to know how many images for final attention pooling
  
      num_images = torch.tensor(
          num_images, device=device, dtype=torch.long
      )
  
      # to patches
  
      x = self.to_patch_embedding(patches)
  
      # factorized 2d absolute positional embedding
  
      h_indices, w_indices = patch_positions.unbind(dim=-1)
  
      h_pos = self.pos_embed_height[h_indices]
      w_pos = self.pos_embed_width[w_indices]
  
      x = x + h_pos + w_pos
  
      # embed dropout
  
      x = self.dropout(x)
  
      # attention
  
      x = self.transformer(x, attn_mask=attn_mask)
  
      # do attention pooling at the end
  
      max_queries = num_images.amax().item()
  
      queries = repeat(
          self.attn_pool_queries,
          "d -> b n d",
          n=max_queries,
          b=x.shape[0],
      )
  
      # attention pool mask
  
      image_id_arange = arange(max_queries)
  
      attn_pool_mask = rearrange(
          image_id_arange, "i -> i 1"
      ) == rearrange(batched_image_ids, "b j -> b 1 j")
  
      attn_pool_mask = attn_pool_mask & rearrange(
          key_pad_mask, "b j -> b 1 j"
      )
  
      attn_pool_mask = rearrange(attn_pool_mask, "b i j -> b 1 i j")
  
      # attention pool
  
      x = (
          self.attn_pool(queries, context=x, attn_mask=attn_pool_mask)
          + queries
      )
  
      x = rearrange(x, "b n d -> (b n) d")
  
      # each batch element may not have same amount of images
  
      is_images = image_id_arange < rearrange(num_images, "b -> b 1")
      is_images = rearrange(is_images, "b n -> (b n)")
  
      x = x[is_images]
  
      # project out to logits
  
      x = self.to_latent(x)
  
      return self.mlp_head(x)
  
imports: from functools import partial
  from typing import List, Union
  import torch
  import torch.nn.functional
  from torch import nn, Tensor
  from torch.nn.utils.rnn import pad_sequence
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 3 for calling the NaViT obj's underlying member method code with example usage: obj.forward(batched_images, group_images, group_max_seq_len) and returns: self