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","`","."]
code: def __init__(
      self,
      pytree_extension_context: Optional[
          _PyTreeExtensionContext
      ] = None,
  ):
      super().__init__()
      self._pytree_extension_context = (
          pytree_extension_context or _PyTreeExtensionContext()
      )
  
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 1 to create a new DynamoFlattenOutputStep object from inherited base classes: io_adapter with example: obj = DynamoFlattenOutputStep(pytree_extension_context)