class_docstr: Flatten nested collection and custom python types and return a flat list of elements.
  
  Extended from :class:`io_adapter.FlattenOutputStep` to support flattening arbitrary
  types via pytree extension. By default this supports many common user defined python
  types such as :class:`ModelOutput` from HuggingFace transformers.
  
  The pytree extension can be customized by passing in a ``_PyTreeExtensionContext``
  object. See :meth:`_PyTreeExtensionContext.register_pytree_node`.
class_docstr_tok: ["Flatten","nested","collection","and","custom","python","types","and","return","a","flat","list","of","elements",".","Extended","from",":","class",":","`","io_adapter.FlattenOutputStep","`","to","support","flattening","arbitrary","types","via","pytree","extension",".","By","default","this","supports","many","common","user","defined","python","types","such","as",":","class",":","`","ModelOutput","`","from","HuggingFace","transformers",".","The","pytree","extension","can","be","customized","by","passing","in","a","``","_PyTreeExtensionContext","``","object",".","See",":","meth",":","`","_PyTreeExtensionContext.register_pytree_node","`","."]
docstr: Flatten the model outputs, under the context of pytree extension.
docstr_tok: ["Flatten","the","model","outputs",",","under","the","context","of","pytree","extension","."]
code: def apply(
      self,
      model_outputs: Any,
      model: Optional[
          Union[torch.nn.Module, Callable, torch_export.ExportedProgram]
      ] = None,
  ) -> Sequence[Any]:
      """Flatten the model outputs, under the context of pytree extension."""
      with self._pytree_extension_context:
          return super().apply(model_outputs, model=model)
  
imports: from __future__ import annotations
  import contextlib
  import functools
  import inspect
  from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
  import torch._dynamo
  import torch.export
  import torch.fx
  import torch.onnx
  from torch.onnx._internal import _beartype, exporter, io_adapter
  from torch.utils import _pytree
label_desc: Use image node_id 2 for calling the DynamoFlattenOutputStep obj's underlying member method code with example usage: obj.apply(model_outputs, model) and returns: super