code: def compute_mask(dims, window_size, shift_size, device):
      """Computing region masks based on: "Liu et al.,
      Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
      <https://arxiv.org/abs/2103.14030>"
      https://github.com/microsoft/Swin-Transformer
  
       Args:
          dims: dimension values.
          window_size: local window size.
          shift_size: shift size.
          device: device.
      """
  
      cnt = 0
  
      if len(dims) == 3:
          d, h, w = dims
          img_mask = torch.zeros((1, d, h, w, 1), device=device)
          for d in (
              slice(-window_size[0]),
              slice(-window_size[0], -shift_size[0]),
              slice(-shift_size[0], None),
          ):
              for h in (
                  slice(-window_size[1]),
                  slice(-window_size[1], -shift_size[1]),
                  slice(-shift_size[1], None),
              ):
                  for w in (
                      slice(-window_size[2]),
                      slice(-window_size[2], -shift_size[2]),
                      slice(-shift_size[2], None),
                  ):
                      img_mask[:, d, h, w, :] = cnt
                      cnt += 1
  
      elif len(dims) == 2:
          h, w = dims
          img_mask = torch.zeros((1, h, w, 1), device=device)
          for h in (
              slice(-window_size[0]),
              slice(-window_size[0], -shift_size[0]),
              slice(-shift_size[0], None),
          ):
              for w in (
                  slice(-window_size[1]),
                  slice(-window_size[1], -shift_size[1]),
                  slice(-shift_size[1], None),
              ):
                  img_mask[:, h, w, :] = cnt
                  cnt += 1
  
      mask_windows = window_partition(img_mask, window_size)
      mask_windows = mask_windows.squeeze(-1)
      attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
      attn_mask = attn_mask.masked_fill(
          attn_mask != 0, float(-100.0)
      ).masked_fill(attn_mask == 0, float(0.0))
  
      return attn_mask
  
imports: from __future__ import annotations
  import itertools
  from collections.abc import Sequence
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch.nn import LayerNorm
  from typing_extensions import Final
  from monai.networks.blocks import MLPBlock
  from monai.networks.blocks import PatchEmbed, UnetOutBlock, UnetrBasicBlock, UnetrUpBlock
  from monai.networks.layers import DropPath, trunc_normal_
  from monai.utils import ensure_tuple_rep, look_up_option, optional_import
  from monai.utils.deprecate_utils import deprecated_arg
label_desc: Use image node_id 4 for calling a global function with example usage: compute_mask(dims, window_size, shift_size, device) and returns: attn_mask