code: def _extract_and_flatten_nested_structure(data, flattened=None):
      
      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
docs: 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.