code: def block_diagonal(mat, block_size):
      """
      Takes a block diagonal tensor of shape (..., B x M, B x N) and returns a tensor
      of shape (..., B, M, N).
  
      :param torch.Tensor mat: an input tensor with 2 or more dimensions
      :param int block_size: the number of blocks B.
      :returns torch.Tensor: a tensor with dimension `mat.dim() + 1`
      """
      B = block_size
      M = mat.size(-2) // B
      N = mat.size(-1) // B
      assert mat.shape[-2:] == (B * M, B * N)
      mat = mat.reshape(mat.shape[:-2] + (B, M, B, N))
      mat = mat.transpose(-2, -3)
      mat = mat.reshape(mat.shape[:-4] + (B * B, M, N))
      return mat[..., :: B + 1, :, :]
  
imports: import math
  import torch
  from torch.fft import irfft, rfft
  from ..None import settings
label_desc: Use image node_id 4 for calling a global function with example usage: block_diagonal(mat, block_size) and returns: mat