mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-07-13 12:05:48 +00:00
393f594051
* [fx] move meta registration * [fx] fix tests. * [fx] fix test. * [fx] fix. * [meta] refactor meta registration.py. * [fx] add compatibility descriptions. * [fx] polish import. * [fx] add a decorator. * [fx] fix tests. * [fx] remove print. * [fx] edit raise error. * [fx] edit raise error. * [fx] add type hint. * [fx] fix import in experimental. * [rpc] remove color debug. * [meta] fix naming.
49 lines
1.1 KiB
Python
49 lines
1.1 KiB
Python
# for PyTorch 1.11 compatibility uses
|
|
from typing import Dict, List, Tuple, Union
|
|
|
|
import torch
|
|
from torch.fx import GraphModule, Node
|
|
|
|
from ..._compatibility import compatibility
|
|
|
|
__all__ = ["calculate_fwd_in", "calculate_fwd_tmp", "calculate_fwd_out"]
|
|
|
|
|
|
@compatibility(is_backward_compatible=True)
|
|
def calculate_fwd_in(n: Node) -> bool:
|
|
"""A helper function to calculate `fwd_in`
|
|
|
|
Args:
|
|
n (Node): a node from the graph
|
|
|
|
Returns:
|
|
save_fwd_in (bool): the result of `save_fwd_in`
|
|
"""
|
|
return n.meta['save_fwd_in']
|
|
|
|
|
|
@compatibility(is_backward_compatible=True)
|
|
def calculate_fwd_tmp(n: Node) -> int:
|
|
"""A helper function to calculate `fwd_tmp`
|
|
|
|
Args:
|
|
n (Node): a node from the graph
|
|
|
|
Returns:
|
|
fwd_tmp (int): the result of `fwd_tmp`
|
|
"""
|
|
return n.meta["fwd_mem_tmp"]
|
|
|
|
|
|
@compatibility(is_backward_compatible=True)
|
|
def calculate_fwd_out(n: Node) -> int:
|
|
"""A helper function to calculate `fwd_out`
|
|
|
|
Args:
|
|
n (Node): a node from the graph
|
|
|
|
Returns:
|
|
fwd_out (int): the result of `fwd_out`
|
|
"""
|
|
return n.meta['fwd_mem_out']
|