(node: torch.fx.Node, node_repr: Union[TensorRepr, TensorReprList])
| 1684 | |
| 1685 | |
| 1686 | def set_node_repr(node: torch.fx.Node, node_repr: Union[TensorRepr, TensorReprList]): |
| 1687 | if isinstance(node_repr, TensorReprList): |
| 1688 | # Convert to a regular list so taht `set_node_spec_attr` can attach each entry |
| 1689 | # to a separate TensorSpec |
| 1690 | node_repr_list = [node_repr[i] for i in range(num_tensors_in_node(node))] |
| 1691 | set_node_spec_attr(node, "etvk_node_repr", node_repr_list) |
| 1692 | else: |
| 1693 | set_node_spec_attr(node, "etvk_node_repr", node_repr) |
| 1694 | |
| 1695 | |
| 1696 | def get_node_repr(node) -> Union[TensorRepr, TensorReprList]: |
nothing calls this directly
no test coverage detected