r""" Return the list of the tensor specs for the node or empty list if the node has no tensor specs.
(
node: torch.fx.Node,
)
| 890 | |
| 891 | |
| 892 | def get_node_tensor_specs( |
| 893 | node: torch.fx.Node, |
| 894 | ) -> Union[List[TensorSpec], Tuple[TensorSpec]]: |
| 895 | r""" |
| 896 | Return the list of the tensor specs for the node or empty list if the node |
| 897 | has no tensor specs. |
| 898 | """ |
| 899 | # get tensor specs |
| 900 | if node.target == memory.view: |
| 901 | base = node.args[0] |
| 902 | assert isinstance(base, torch.fx.Node) |
| 903 | specs = base.meta.get("spec") |
| 904 | else: |
| 905 | specs = node.meta.get("spec") |
| 906 | |
| 907 | if isinstance(specs, TensorSpec): |
| 908 | specs = [specs] |
| 909 | if not isinstance(specs, (list, tuple)): |
| 910 | return [] |
| 911 | else: |
| 912 | return [ |
| 913 | spec |
| 914 | for spec in specs |
| 915 | if not isinstance(spec, (int, float, bool, str, type(None))) |
| 916 | ] |
| 917 | |
| 918 | |
| 919 | # Little bit hacky to check if the graph contains |