Flatten the structure of a nested container to a list.

    Parameters
    ----------
    data : A single NDArray/Symbol or nested container with NDArrays/Symbol.
        The nested container to be flattened.
    flattened : list or None
        The container thats holds flattened result.
    Returns
    -------
    structure : An integer or a nested container with integers.
        The extracted structure of the container of `data`.
    flattened : (optional) list
        The container thats holds flattened result.
        It is returned only when the input argument `flattened` is not given.

def _extract_and_flatten_nested_structure(data, flattened=None):
    """Flatten the structure of a nested container to a list.

    Parameters
    ----------
    data : A single NDArray/Symbol or nested container with NDArrays/Symbol.
        The nested container to be flattened.
    flattened : list or None
        The container thats holds flattened result.
    Returns
    -------
    structure : An integer or a nested container with integers.
        The extracted structure of the container of `data`.
    flattened : (optional) list
        The container thats holds flattened result.
        It is returned only when the input argument `flattened` is not given.
    """
    if flattened is None:
        flattened = []
        structure = _extract_and_flatten_nested_structure(data, flattened)
        return structure, flattened
    if isinstance(data, list):
        return list(_extract_and_flatten_nested_structure(x, flattened) for x in data)
    elif isinstance(data, tuple):
        return tuple(_extract_and_flatten_nested_structure(x, flattened) for x in data)
    elif isinstance(data, dict):
        return {k: _extract_and_flatten_nested_structure(v) for k, v in data.items()}
    elif isinstance(data, (mx.sym.Symbol, mx.nd.NDArray)):
        flattened.append(data)
        return len(flattened) - 1
    else:
        raise NotImplementedError