class_docstr: Vision transformer as per https://arxiv.org/abs/2010.11929.
class_docstr_tok: ["Vision","transformer","as","per","https",":","\/\/arxiv.org\/abs\/2010.11929","."]
docstr: Computes model complexity. If you alter the model, make sure to update.
code: def complexity(cx, params=None):
      """Computes model complexity. If you alter the model, make sure to update."""
      p = ViT.get_params() if not params else params
      ViT.check_params(p)
      if p["stem_type"] == "patchify":
          cx = ViTStemPatchify.complexity(
              cx, 3, p["hidden_d"], p["patch_size"]
          )
      elif p["stem_type"] == "conv":
          ks, ws, ss = (
              p["c_stem_kernels"],
              p["c_stem_dims"],
              p["c_stem_strides"],
          )
          cx = ViTStemConv.complexity(cx, 3, ks, ws, ss)
      seq_len = (p["image_size"] // cfg.VIT.PATCH_SIZE) ** 2
      if p["cls_type"] == "token":
          seq_len += 1
          cx["params"] += p["hidden_d"]
      # Params of position embeddings
      cx["params"] += seq_len * p["hidden_d"]
      cx = ViTEncoder.complexity(
          cx,
          p["n_layers"],
          p["hidden_d"],
          p["n_heads"],
          p["mlp_d"],
          seq_len,
      )
      cx = ViTHead.complexity(cx, p["hidden_d"], p["num_classes"])
      return cx
  
imports: import math
  import numpy
  import torch
  from pycls.core.config import cfg
  from pycls.models.blocks import activation, conv2d, conv2d_cx, layernorm, layernorm_cx, linear, linear_cx, MultiheadAttention, norm2d, norm2d_cx, patchify2d, patchify2d_cx
  from torch.nn import init, Module, Parameter
label_desc: Use image node_id 5 for calling the ViT obj's underlying member method code with example usage: obj.complexity(cx, params) and returns: cx