code: def rearrange_heads_inner(t: torch.Tensor, h: int) -> torch.Tensor:
      b = t.shape[0] // h
      n = t.shape[1]
      return t.view(b, h, n, -1).transpose(1, 2).reshape(b, n, -1)
  
imports: import math
  from inspect import isfunction
  import torch
  import torch.nn.functional
  from apex.contrib.group_norm import GroupNorm
  from einops import rearrange
  from torch import einsum, nn
  from torch._dynamo import disable
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import checkpoint
label_desc: Use image node_id 10 for calling a global function with example usage: rearrange_heads_inner(t, h) and returns: t