code: def lmcl_matmul(
      i: torch.Tensor,
      w: torch.Tensor,
      tgt: torch.Tensor,
      w_idx: int,
      margin: float,
      scale: Optional[float],
  ) -> torch.Tensor:
      """LMCL variation of matmul with normalization, margin and scale."""
      # normalize and matmul
      logits = torch.matmul(
          F.normalize(i, dim=1), F.normalize(w, dim=1).T
      )
  
      # add margin using a mask since tgt might be out of the the weight split's range.
      mask = torch.arange(
          w_idx * w.shape[0],
          (w_idx + 1) * w.shape[0],
          dtype=torch.long,
          device=i.device,
      ).expand(i.shape[0], -1)
      logits[mask == tgt.reshape(-1, 1)] -= margin
  
      # add scale
      logits *= scale
  
      return logits
  
imports: from typing import Any, Optional, Tuple
  import torch
  from torch import nn
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 4 for calling a global function with example usage: lmcl_matmul(i, w, tgt, w_idx, margin, scale) and returns: logits