class_docstr: Context manager to register PyTree extension.
class_docstr_tok: ["Context","manager","to","register","PyTree","extension","."]
docstr: Register PyTree extension for a custom python type.
  
  Args:
      class_type: The custom python type.
      flatten_func: The flatten function.
      unflatten_func: The unflatten function.
  
  Raises:
      AssertionError: If the custom python type is already registered.
docstr_tok: ["Register","PyTree","extension","for","a","custom","python","type",".","Args",":","class_type",":","The","custom","python","type",".","flatten_func",":","The","flatten","function",".","unflatten_func",":","The","unflatten","function",".","Raises",":","AssertionError",":","If","the","custom","python","type","is","already","registered","."]
code: def register_pytree_node(
      self,
      class_type: Type,
      flatten_func: pytree.FlattenFunc,
      unflatten_func: pytree.UnflattenFunc,
  ):
      """Register PyTree extension for a custom python type.
  
      Args:
          class_type: The custom python type.
          flatten_func: The flatten function.
          unflatten_func: The unflatten function.
  
      Raises:
          AssertionError: If the custom python type is already registered.
      """
      if (
          class_type in pytree.SUPPORTED_NODES
          or class_type in self._extensions
      ):
          # PyTree node already registered.
          # E.g., `huggingface/transformer` registers `ModelOutput` as PyTree node after
          # https://github.com/huggingface/transformers/pull/25358.
          return
      self._extensions[class_type] = (flatten_func, unflatten_func)
  
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 4 for calling the _PyTreeExtensionContext obj's underlying member method code with example usage: obj.register_pytree_node(class_type, flatten_func, unflatten_func) without return types