mirror of https://github.com/hpcaitech/ColossalAI
29 lines
1004 B
Python
29 lines
1004 B
Python
from .api import (
|
|
compute_global_numel,
|
|
customized_distributed_tensor_to_param,
|
|
distribute_tensor,
|
|
distribute_tensor_with_customization,
|
|
get_device_mesh,
|
|
get_global_shape,
|
|
get_layout,
|
|
get_sharding_spec,
|
|
is_customized_distributed_tensor,
|
|
is_distributed_tensor,
|
|
is_sharded,
|
|
redistribute,
|
|
shard_colwise,
|
|
shard_rowwise,
|
|
sharded_tensor_to_param,
|
|
to_global,
|
|
to_global_for_customized_distributed_tensor,
|
|
)
|
|
from .layout import Layout
|
|
from .sharding_spec import ShardingSpec
|
|
|
|
__all__ = [
|
|
'is_distributed_tensor', 'distribute_tensor', 'to_global', 'is_sharded', 'shard_rowwise', 'shard_colwise',
|
|
'sharded_tensor_to_param', 'compute_global_numel', 'get_sharding_spec', 'get_global_shape', 'get_device_mesh',
|
|
'redistribute', 'get_layout', 'is_customized_distributed_tensor', 'distribute_tensor_with_customization',
|
|
'to_global_for_customized_distributed_tensor', 'customized_distributed_tensor_to_param', 'Layout', 'ShardingSpec'
|
|
]
|