class_docstr: Transposed LayerNorm module.
  
  LayerNorm module over second dimension
  Assumes shape (B, D, L), where L can be 1 or more axis
  
  This is slow and a dedicated CUDA/Triton implementation
  shuld provide substantial end-to-end speedup
code: def __init__(self, d, scalar=True):
      super().__init__()
      self.scalar = scalar
      if self.scalar:
          self.m = nn.Parameter(torch.zeros(1))
          self.s = nn.Parameter(torch.ones(1))
          setattr(self.m, "_optim", {"weight_decay": 0.0})
          setattr(self.s, "_optim", {"weight_decay": 0.0})
      else:
          self.ln = nn.LayerNorm(d)
  
imports: import math
  from functools import partial
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from opt_einsum import contract
label_desc: Use image node_id 1 to create a new TransposedLN object from inherited base classes: nn with example: obj = TransposedLN(d, scalar)