ColossalAI/colossalai/legacy/nn/layer/utils/__init__.py

22 lines
445 B
Python
Raw Normal View History

from .common import (
ACT2FN,
CheckpointModule,
_ntuple,
divide,
get_tensor_parallel_mode,
set_tensor_parallel_attribute_by_partition,
set_tensor_parallel_attribute_by_size,
to_2tuple,
)
__all__ = [
"CheckpointModule",
"divide",
"ACT2FN",
"set_tensor_parallel_attribute_by_size",
"set_tensor_parallel_attribute_by_partition",
"get_tensor_parallel_mode",
"_ntuple",
"to_2tuple",
]