This methods reshapes all values in a dictionary.

    The indices from start to stop will be flattened into a single index.

    Args:
        weights: A dictionary mapping keys to numpy arrays.
        start: The starting index.
        stop: The ending index.

def flatten(weights, start=0, stop=2):
    """This methods reshapes all values in a dictionary.

    The indices from start to stop will be flattened into a single index.

    Args:
        weights: A dictionary mapping keys to numpy arrays.
        start: The starting index.
        stop: The ending index.
    """
    for key, val in weights.items():
        new_shape = val.shape[0:start] + (-1, ) + val.shape[stop:]
        weights[key] = val.reshape(new_shape)
    return weights