github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 18:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx]refactor tracer (#1335)

Browse Source
This commit is contained in:
YuliangLiu0306
2022-07-19 15:50:42 +08:00
committed by GitHub
parent bf5066fba7
commit 4631fef8a0
8 changed files with 50 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
colossalai/fx/tracer/_tracer_utils.py
View File

@@ -1,5 +1,5 @@
from typing import List, Union, Any
from ..proxy import ColoProxy, MetaDeviceAttribute
from ..proxy import ColoProxy, ColoAttribute
__all__ = ['is_element_in_list', 'extract_meta']
@@ -19,10 +19,11 @@ def is_element_in_list(elements: Union[List[Any], Any], list_: List[Any]):
def extract_meta(*args, **kwargs):
def _convert(val):
if isinstance(val, MetaDeviceAttribute):
return 'meta'
elif isinstance(val, ColoProxy):
if isinstance(val, ColoProxy):
return val.meta_data
elif isinstance(val, (list, tuple)):
return type(val)([_convert(ele) for ele in val])
return val
new_args = [_convert(val) for val in args]