mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
7982b8c703
Summary: The definition of `ExternKernelNode` and `ExternKernelNodes` schema in `torch/_export/serde/aoti_schema.py` is a complete duplicate of the ones in `torch/_export/serde/schema.py`. Test Plan: CI Rollback Plan: Differential Revision: D76558294 Pull Request resolved: https://github.com/pytorch/pytorch/pull/155867 Approved by: https://github.com/jingsh
25 lines
830 B
Python
25 lines
830 B
Python
import json
|
|
|
|
from torch._export.serde.schema import ExternKernelNode, ExternKernelNodes, Node
|
|
from torch._export.serde.serialize import _dataclass_to_dict, EnumEncoder
|
|
from torch._inductor.ir import ExternKernelNode as inductor_ExternKernelNode
|
|
|
|
|
|
def serialize_extern_kernel_node(
|
|
extern_kernel_node: inductor_ExternKernelNode,
|
|
) -> ExternKernelNode:
|
|
assert isinstance(extern_kernel_node.node, Node)
|
|
return ExternKernelNode(
|
|
name=extern_kernel_node.name,
|
|
node=extern_kernel_node.node,
|
|
)
|
|
|
|
|
|
def extern_node_json_serializer(
|
|
extern_kernel_nodes: list[inductor_ExternKernelNode],
|
|
) -> str:
|
|
serialized_nodes = ExternKernelNodes(
|
|
nodes=[serialize_extern_kernel_node(node) for node in extern_kernel_nodes]
|
|
)
|
|
return json.dumps(_dataclass_to_dict(serialized_nodes), cls=EnumEncoder)
|