mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-08-25 19:21:17 +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.
45 lines
782 B
Python
45 lines
782 B
Python
from operator import add, floordiv, getitem, mul, neg, pos, setitem, sub
|
|
|
|
import torch
|
|
|
|
__all__ = ['INPLACE_OPS', 'INPLACE_METHOD', 'NON_INPLACE_METHOD']
|
|
|
|
# TODO fill out the inplace ops
|
|
INPLACE_OPS = [
|
|
add,
|
|
sub,
|
|
mul,
|
|
floordiv,
|
|
neg,
|
|
pos,
|
|
getitem,
|
|
setitem,
|
|
getattr,
|
|
torch.Tensor.cpu,
|
|
]
|
|
|
|
# TODO: list all call_methods that are inplace here
|
|
INPLACE_METHOD = [
|
|
'transpose',
|
|
'permute',
|
|
# TODO: reshape may return a copy of the data if the data is not contiguous
|
|
'reshape',
|
|
'dim',
|
|
'flatten',
|
|
'size',
|
|
'view',
|
|
'unsqueeze',
|
|
'to',
|
|
'type',
|
|
'flatten',
|
|
]
|
|
|
|
# TODO: list all call_methods that are not inplace here
|
|
NON_INPLACE_METHOD = [
|
|
'chunk',
|
|
'contiguous',
|
|
'expand',
|
|
'mean',
|
|
'split',
|
|
]
|