def MergeHeads(n_heads, merged_batch_and_head=True):
  """Returns a layer that rejoins heads, after multi-head computation."""
  def f(x):
    if merged_batch_and_head:
      dim_0, seq_len, d_head = x.shape
      if dim_0 % n_heads != 0:
        raise ValueError(
            f"Array's leading dimension ({dim_0}) is not a multiple of the"
            f" number of attention heads ({n_heads}).")

      batch_size = dim_0 // n_heads
      x = x.reshape((batch_size, n_heads, seq_len, d_head))
    else:
      batch_size, _, seq_len, d_head = x.shape

    # (b_size, n_heads, seq_len, d_head) --> (b_size, seq_len, d_feature)
    x = x.transpose((0, 2, 1, 3))
    x = x.reshape((batch_size, seq_len, n_heads * d_head))
    return x
  return Fn('MergeHeads', f)
Returns a layer that rejoins heads, after multi-head computation.
trax/layers/attention.py