def gather_array(array):
    """
    Gathers an array that has previously been distributed using `distribute_array`,
    by removing the leading axis and stacking them on the second axis.
    """
    s = array.shape
    assert len(s) > 2, "Array must have at least 3 axes"
    return array.reshape(s[0] + s[1], s[2:])