code: def register_pytree_node(
      cls: Type[Any],
      flatten_fn: FlattenFunc,
      unflatten_fn: UnflattenFunc,
      *,
      serialized_type_name: Optional[str] = None,
      to_dumpable_context: Optional[ToDumpableContextFn] = None,
      from_dumpable_context: Optional[FromDumpableContextFn] = None,
  ) -> None:
      """Register a container-like type as pytree node.
  
      Args:
          cls (type): A Python type to treat as an internal pytree node.
          flatten_fn (callable): A function to be used during flattening, taking an instance of
              ``cls`` and returning a pair, with (1) an iterable for the children to be flattened
              recursively, and (2) some hashable auxiliary data to be stored in the treespec and to be
              passed to the ``unflatten_fn``.
          unflatten_fn (callable): A function taking two arguments: the auxiliary data that was
              returned by ``flatten_fn`` and stored in the treespec, and the unflattened children.
              The function should return an instance of ``cls``.
          serialized_type_name (str, optional): A keyword argument used to specify the fully
              qualified name used when serializing the tree spec.
          to_dumpable_context (callable, optional): An optional keyword argument to custom specify how
              to convert the context of the pytree to a custom json dumpable representation. This is
              used for json serialization, which is being used in :mod:`torch.export` right now.
          from_dumpable_context (callable, optional): An optional keyword argument to custom specify
              how to convert the custom json dumpable representation of the context back to the
              original context. This is used for json deserialization, which is being used in
              :mod:`torch.export` right now.
  
      Example::
  
          >>> # xdoctest: +SKIP
          >>> # Registry a Python type with lambda functions
          >>> register_pytree_node(
          ...     set,
          ...     lambda s: (sorted(s), None, None),
          ...     lambda children, _: set(children),
          ... )
      """
      _private_register_pytree_node(
          cls,
          flatten_fn,
          unflatten_fn,
          serialized_type_name=serialized_type_name,
          to_dumpable_context=to_dumpable_context,
          from_dumpable_context=from_dumpable_context,
      )
  
      from . import _pytree as python
  
      python._private_register_pytree_node(
          cls,
          flatten_fn,
          unflatten_fn,
          serialized_type_name=serialized_type_name,
          to_dumpable_context=to_dumpable_context,
          from_dumpable_context=from_dumpable_context,
      )
  
imports: import functools
  import warnings
  from typing import Any, Callable, Iterable, List, Optional, overload, Tuple, Type, TypeVar, Union
  import torch
  import optree
  from optree import PyTreeSpec
label_desc: Use image node_id 2 for calling a global function with example usage: register_pytree_node(cls, flatten_fn, unflatten_fn) without return types