code: def split_heads(x: mx.sym.Symbol, depth_per_head: int, heads: int) -> mx.sym.Symbol:
      
      # (batch, length, heads, depth_per_head)
      x = mx.sym.reshape(data=x, shape=(0, -1, heads, depth_per_head))
      # (batch, heads, length, depth/heads)
      x = mx.sym.transpose(data=x, axes=(0, 2, 1, 3))
      # (batch * heads, length, depth/heads)
      return mx.sym.reshape(data=x, shape=(-3, -1, depth_per_head))
docs: Returns a symbol with head dimension folded into batch and depth divided by the number of heads.
  
      :param x: Symbol of shape (batch, length, depth).
      :param depth_per_head: Depth per head.
      :param heads: Number of heads.
      :return: Symbol of shape (batch * heads, length, depth_per_heads).