_register_dataclass_output_as_pytree() — pytorch Function Reference
Architecture documentation for the _register_dataclass_output_as_pytree() function in common.py from the pytorch codebase.
Entity Profile
Dependency Diagram
graph TD 8964c48d_7b16_2455_878e_dbe471c4a037["_register_dataclass_output_as_pytree()"] 3df8de63_d71a_714c_93a1_1a0e24c0f362["load()"] 3df8de63_d71a_714c_93a1_1a0e24c0f362 -->|calls| 8964c48d_7b16_2455_878e_dbe471c4a037 d4bae6c8_efa9_4b5e_53ab_4e403ff2693e["load()"] d4bae6c8_efa9_4b5e_53ab_4e403ff2693e -->|calls| 8964c48d_7b16_2455_878e_dbe471c4a037 f258a26f_4fc8_d6ab_1167_c608fca925fb["export()"] f258a26f_4fc8_d6ab_1167_c608fca925fb -->|calls| 8964c48d_7b16_2455_878e_dbe471c4a037 style 8964c48d_7b16_2455_878e_dbe471c4a037 fill:#6366f1,stroke:#818cf8,color:#fff
Relationship Graph
Source Code
benchmarks/dynamo/common.py lines 766–780
def _register_dataclass_output_as_pytree(example_outputs) -> None:
# NOTE(angelayi): For huggingface benchmark, some example outputs are
# formatted as a dataclass which pytree cannot consume. So we want
# to register the pytree implementation here
example_outputs_flat = pytree.tree_leaves(example_outputs)
output_dataclass_types = [
type(out) for out in example_outputs_flat if dataclasses.is_dataclass(type(out))
]
for output_type in output_dataclass_types:
from torch._export.utils import register_dataclass_as_pytree_node
register_dataclass_as_pytree_node(
output_type,
serialized_type_name=f"{output_type.__module__}.{output_type.__name__}",
)Domain
Subdomains
Source
Frequently Asked Questions
What does _register_dataclass_output_as_pytree() do?
_register_dataclass_output_as_pytree() is a function in the pytorch codebase.
What calls _register_dataclass_output_as_pytree()?
_register_dataclass_output_as_pytree() is called by 3 function(s): export, load, load.
Analyze Your Own Codebase
Get architecture documentation, dependency graphs, and domain analysis for your codebase in minutes.
Try Supermodel Free