class_docstr: Implements a PARSeq architecture as described in `"Scene Text Recognition
  with Permuted Autoregressive Sequence Models" <https://arxiv.org/pdf/2207.06966>`_.
  Slightly modified implementation based on the official Pytorch implementation: <https://github.com/baudm/parseq/tree/main`_.
  
  Args:
  ----
      feature_extractor: the backbone serving as feature extractor
      vocab: vocabulary used for encoding
      embedding_units: number of embedding units
      max_length: maximum word length handled by the model
      dropout_prob: dropout probability for the decoder
      dec_num_heads: number of attention heads in the decoder
      dec_ff_dim: dimension of the feed forward layer in the decoder
      dec_ffd_ratio: depth multiplier for the feed forward layer in the decoder
      input_shape: input shape of the image
      exportable: onnx exportable returns only logits
      cfg: dictionary containing information about the model
class_docstr_tok: ["Implements","a","PARSeq","architecture","as","described","in","`","``","Scene","Text","Recognition","with","Permuted","Autoregressive","Sequence","Models","''","<","https",":","\/\/arxiv.org\/pdf\/2207.06966",">","`","_",".","Slightly","modified","implementation","based","on","the","official","Pytorch","implementation",":","<","https",":","\/\/github.com\/baudm\/parseq\/tree\/main","`","_",".","Args",":","--","--","feature_extractor",":","the","backbone","serving","as","feature","extractor","vocab",":","vocabulary","used","for","encoding","embedding_units",":","number","of","embedding","units","max_length",":","maximum","word","length","handled","by","the","model","dropout_prob",":","dropout","probability","for","the","decoder","dec_num_heads",":","number","of","attention","heads","in","the","decoder","dec_ff_dim",":","dimension","of","the","feed","forward","layer","in","the","decoder","dec_ffd_ratio",":","depth","multiplier","for","the","feed","forward","layer","in","the","decoder","input_shape",":","input","shape","of","the","image","exportable",":","onnx","exportable","returns","only","logits","cfg",":","dictionary","containing","information","about","the","model"]
code: def forward(
      self,
      x: torch.Tensor,
      target: Optional[List[str]] = None,
      return_model_output: bool = False,
      return_preds: bool = False,
  ) -> Dict[str, Any]:
      features = self.feat_extractor(x)[
          "features"
      ]  # (batch_size, patches_seqlen, d_model)
      # remove cls token
      features = features[:, 1:, :]
  
      if self.training and target is None:
          raise ValueError("Need to provide labels during training")
  
      if target is not None:
          # Build target tensor
          _gt, _seq_len = self.build_target(target)
          gt, seq_len = torch.from_numpy(_gt).to(dtype=torch.long).to(
              x.device
          ), torch.tensor(_seq_len).to(x.device)
          gt = gt[
              :, : int(seq_len.max().item()) + 2
          ]  # slice up to the max length of the batch + 2 (SOS + EOS)
  
          if self.training:
              # Generate permutations for the target sequences
              tgt_perms = self.generate_permutations(seq_len)
  
              gt_in = gt[
                  :, :-1
              ]  # remove EOS token from longest target sequence
              gt_out = gt[:, 1:]  # remove SOS token
              # Create padding mask for target input
              # [True, True, True, ..., False, False, False] -> False is masked
              padding_mask = (
                  ~(
                      (
                          (gt_in == self.vocab_size + 2)
                          | (gt_in == self.vocab_size)
                      )
                      .int()
                      .cumsum(-1)
                      > 0
                  )
                  .unsqueeze(1)
                  .unsqueeze(1)
              )  # (N, 1, 1, seq_len)
  
              loss = torch.tensor(0.0, device=features.device)
              loss_numel: Union[int, float] = 0
              n = (gt_out != self.vocab_size + 2).sum().item()
              for i, perm in enumerate(tgt_perms):
                  (
                      _,
                      target_mask,
                  ) = self.generate_permutations_attention_masks(
                      perm
                  )  # (seq_len, seq_len)
                  # combine both masks
                  mask = (
                      target_mask.bool() & padding_mask.bool()
                  ).int()  # (N, 1, seq_len, seq_len)
  
                  logits = self.head(
                      self.decode(gt_in, features, mask)
                  ).flatten(end_dim=1)
                  loss += n * F.cross_entropy(
                      logits,
                      gt_out.flatten(),
                      ignore_index=self.vocab_size + 2,
                  )
                  loss_numel += n
                  # After the second iteration (i.e. done with canonical and reverse orderings),
                  # remove the [EOS] tokens for the succeeding perms
                  if i == 1:
                      gt_out = torch.where(
                          gt_out == self.vocab_size,
                          self.vocab_size + 2,
                          gt_out,
                      )
                      n = (gt_out != self.vocab_size + 2).sum().item()
  
              loss /= loss_numel
  
          else:
              gt = gt[:, 1:]  # remove SOS token
              max_len = gt.shape[1] - 1  # exclude EOS token
              logits = self.decode_autoregressive(features, max_len)
              loss = F.cross_entropy(
                  logits.flatten(end_dim=1),
                  gt.flatten(),
                  ignore_index=self.vocab_size + 2,
              )
      else:
          logits = self.decode_autoregressive(features)
  
      logits = _bf16_to_float32(logits)
  
      out: Dict[str, Any] = {}
      if self.exportable:
          out["logits"] = logits
          return out
  
      if return_model_output:
          out["out_map"] = logits
  
      if target is None or return_preds:
          # Post-process boxes
          out["preds"] = self.postprocessor(logits)
  
      if target is not None:
          out["loss"] = loss
  
      return out
  
imports: import math
  from copy import deepcopy
  from itertools import permutations
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  import numpy
  import torch
  from torch import nn
  from torch.nn import functional
  from torchvision.models._utils import IntermediateLayerGetter
  from doctr.datasets import VOCABS
  from doctr.models.modules.transformer import MultiHeadAttention, PositionwiseFeedForward
  from ...classification import vit_s
  from ...utils.pytorch import _bf16_to_float32, load_pretrained_params
  from .base import _PARSeq, _PARSeqPostProcessor
label_desc: Use image node_id 6 for calling the PARSeq obj's underlying member method code with example usage: obj.forward(x, target, return_model_output, return_preds) and returns: out, out