code: def _test_shift_right():
      x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
      inspect(x)
      inspect(shift_right(x))
  
      x = torch.arange(1, 6)[None, :, None, None].repeat(5, 1, 1, 1)
      inspect(x[:, :, 0, 0])
      inspect(shift_right(x)[:, :, 0, 0])
  
      x = torch.arange(1, 6)[None, :, None, None].repeat(3, 1, 1, 1)
      inspect(x[:, :, 0, 0])
      inspect(shift_right(x)[:, :, 0, 0])
  
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 2 for calling a global function with example usage: _test_shift_right() without return types