def SplitIntoHeads(n_heads, merged_batch_and_head=True):
  """Returns a layer that reshapes an array for multi-head computation."""
  def f(x):
    batch_size, seq_len, d_feature = x.shape
    if d_feature % n_heads != 0:
      raise ValueError(
          f'Feature embedding dimensionality ({d_feature}) is not a multiple'
          f' of the requested number of attention heads ({n_heads}).')

    d_head = d_feature // n_heads

    # (b_size, seq_len, d_feature) --> (b_size*n_heads, seq_len, d_head)
    x = x.reshape((batch_size, seq_len, n_heads, d_head))
    x = x.transpose((0, 2, 1, 3))
    if merged_batch_and_head:
      x = x.reshape((batch_size * n_heads, seq_len, d_head))
    return x
  return Fn('SplitIntoHeads', f)
Returns a layer that reshapes an array for multi-head computation.
trax/layers/attention.py