| 1632 | |
| 1633 | |
| 1634 | def set_node_spec_attr(node: torch.fx.Node, attr: str, value): |
| 1635 | assert "spec" in node.meta |
| 1636 | spec = node.meta["spec"] |
| 1637 | if isinstance(spec, TensorSpec): |
| 1638 | setattr(spec, attr, value) |
| 1639 | elif isinstance(spec, (list, tuple)): |
| 1640 | # Special case if value is a list/tuple of the same length as the |
| 1641 | # collection of tensors in the node. In this case, treat the value list |
| 1642 | # as a list of values to set indivudually for each tensor in the node |
| 1643 | if isinstance(value, (list, tuple)) and len(spec) == len(value): |
| 1644 | assert len(spec) == len(value) |
| 1645 | for s, v in zip(spec, value): |
| 1646 | assert isinstance(s, TensorSpec) |
| 1647 | setattr(s, attr, v) |
| 1648 | # Otherwise, set the attribute to value for all tensors in the list |
| 1649 | else: |
| 1650 | for s in spec: |
| 1651 | assert isinstance(s, TensorSpec) |
| 1652 | setattr(s, attr, value) |
| 1653 | else: |
| 1654 | raise RuntimeError(f"Cannot set attr for spec of type {type(spec)}") |
| 1655 | |
| 1656 | |
| 1657 | def get_node_spec_attr(node: torch.fx.Node, attr: str, return_first: bool = True): |