code: def broadcast_to_heads(x: mx.sym.Symbol, num_heads: int, ndim: int, fold_heads: bool = True) -> mx.sym.Symbol:
      
      dims = [0] * (ndim - 1)
      # x: (batch, 1)
      x = mx.sym.expand_dims(x, axis=1)
      # x: (batch, heads, dims...)
      x = mx.sym.broadcast_to(x, shape=[0, num_heads] + dims)
      if fold_heads:
          # (batch * heads, dims...)
          return mx.sym.reshape(x, shape=[-3] + dims)
      else:
          # x: (batch, heads, dims...)
          return x
docs: Broadcasts batch-major input of shape (batch, d1 ... dn-1) to (batch*heads, d1 ... dn-1).
  
      :param x: Batch-major input. Shape: (batch, d1 ... dn-1).
      :param num_heads: Number of heads.
      :param ndim: Number of dimensions in x.
      :param fold_heads: Whether to fold heads dimension into batch dimension.
      :return: Tensor with each sample repeated heads-many times.
               Shape: (batch * heads, d1 ... dn-1) if fold_heads == True, (batch, heads, d1 ... dn-1) else.