code: def shift_right(x: torch.Tensor):
      """
      This method shifts $i^{th}$ row of a matrix by $i$ columns.
  
      If the input is `[[1, 2 ,3], [4, 5 ,6], [7, 8, 9]]`, the shifted
      result would be `[[1, 2 ,3], [0, 4, 5], [6, 0, 7]]`.
      *Ideally we should mask out the lower triangle but it's ok for our purpose*.
      """
  
      # Concatenate a column of zeros
      zero_pad = x.new_zeros(x.shape[0], 1, *x.shape[2:])
      x_padded = torch.cat([x, zero_pad], dim=1)
  
      # Reshape and remove excess elements from the end
      x_padded = x_padded.view(x.shape[1] + 1, x.shape[0], *x.shape[2:])
      x = x_padded[:-1].view_as(x)
  
      #
      return x
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/xl/relative_mha.py
imports: import torch
  from torch import nn
  from labml.logger import inspect
  from labml_nn.transformers.mha import MultiHeadAttention
label_desc: Use image node_id 1 for calling a global function with example usage: shift_right(x) and returns: x