ColossalAI/colossalai/legacy/tensor/__init__.py

18 lines
418 B
Python

from . import distspec
from .compute_spec import ComputePattern, ComputeSpec
from .dist_spec_mgr import DistSpecManager
from .distspec import ReplicaSpec, ShardSpec
from .process_group import ProcessGroup
from .tensor_spec import ColoTensorSpec
__all__ = [
'ComputePattern',
'ComputeSpec',
'distspec',
'DistSpecManager',
'ProcessGroup',
'ColoTensorSpec',
'ShardSpec',
'ReplicaSpec',
]