class_docstr: Context manager to register PyTree extension.
class_docstr_tok: ["Context","manager","to","register","PyTree","extension","."]
code: def __init__(self):
      self._extensions = {}
      # Register PyTree extension for HuggingFace model output.
      self._register_huggingface_model_output_extension()
  
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 _PyTreeExtensionContext object  with example: obj = _PyTreeExtensionContext()