code: def split_heads(
      F, x: Tensor, dim_per_head: int, heads: int
  ) -> Tensor:
      r"""
      Returns a tensor with head dimension folded into batch and last dimension
      divided by the number of heads.
  
      Parameters
      ----------
      x
          Tensor of shape (batch_size, time_length, dim).
      dim_per_head
          Dimension per head
      heads
          Number of heads
  
      Returns
      -------
      Tensor of shape (batch_size * heads, time_length, dim_per_head).
      """
  
      # (batch_size, time_length, heads, dim_per_head)
      x = F.reshape(data=x, shape=(0, -1, heads, dim_per_head))
      # (batch_size, heads, time_length, dim/heads)
      x = F.transpose(data=x, axes=(0, 2, 1, 3))
      # (batch_size * heads, time_length, dim/heads)
      return F.reshape(data=x, shape=(-3, -1, dim_per_head))
  
imports: from typing import Dict, Optional, Tuple
  import mxnet
  from mxnet.gluon import HybridBlock
  from gluonts.mx import Tensor
label_desc: Use image node_id 1 for calling a global function with example usage: split_heads(F, x, dim_per_head, heads) and returns: F