code: def group_images_by_max_seq_len(
      images: List[Tensor],
      patch_size: int,
      calc_token_dropout=None,
      max_seq_len=2048,
  ) -> List[List[Tensor]]:
      calc_token_dropout = default(calc_token_dropout, always(0.0))
  
      groups = []
      group = []
      seq_len = 0
  
      if isinstance(calc_token_dropout, (float, int)):
          calc_token_dropout = always(calc_token_dropout)
  
      for image in images:
          assert isinstance(image, Tensor)
  
          image_dims = image.shape[-2:]
          ph, pw = map(lambda t: t // patch_size, image_dims)
  
          image_seq_len = ph * pw
          image_seq_len = int(
              image_seq_len * (1 - calc_token_dropout(*image_dims))
          )
  
          assert (
              image_seq_len <= max_seq_len
          ), f"image with dimensions {image_dims} exceeds maximum sequence length"
  
          if (seq_len + image_seq_len) > max_seq_len:
              groups.append(group)
              group = []
              seq_len = 0
  
          group.append(image)
          seq_len += image_seq_len
  
      if len(group) > 0:
          groups.append(group)
  
      return groups
  
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 6 for calling a global function with example usage: group_images_by_max_seq_len(images, patch_size, calc_token_dropout, max_seq_len) and returns: groups