File: fx_utils.py

package info (click to toggle)
python-opt-einsum-fx 0.1.4-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 180 kB
  • sloc: python: 664; makefile: 13
file content (14 lines) | stat: -rw-r--r-- 310 bytes parent folder | download 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
 
from typing import Optional

import torch
from torch import fx


def get_shape(n: fx.Node) -> Optional[torch.Size]:
    """Get the shape of a node after ``ShapeProp``"""
    try:
        return n.meta["tensor_meta"].shape
    except KeyError:
        return None
    except AttributeError:
        return None