def expand_shape(shape):
    """
    Expands a flat shape to an expanded shape

    >>> expand_shape([2, 3])
    [2, [3, 3]]
    >>> expand_shape([2, 3, 4])
    [2, [3, 3], [[4, 4, 4], [4, 4, 4]]]
    """
    expanded = [shape[0]]
    for i in range(1, len(shape)):
        next = [shape[i]] * shape[i-1]
        for j in range(1, i):
            next = [next] * shape[j]
        expanded.append(next)
    return expanded