2023-06-22 03:42:11 +00:00
|
|
|
from .api import (
|
|
|
|
|
compute_global_numel,
|
2023-06-23 08:07:09 +00:00
|
|
|
customized_distributed_tensor_to_param,
|
2023-06-22 03:42:11 +00:00
|
|
|
distribute_tensor,
|
2023-06-23 08:07:09 +00:00
|
|
|
distribute_tensor_with_customization,
|
2023-06-22 03:42:11 +00:00
|
|
|
get_device_mesh,
|
|
|
|
|
get_global_shape,
|
|
|
|
|
get_layout,
|
|
|
|
|
get_sharding_spec,
|
2023-06-23 08:07:09 +00:00
|
|
|
is_customized_distributed_tensor,
|
2023-06-22 03:42:11 +00:00
|
|
|
is_distributed_tensor,
|
|
|
|
|
is_sharded,
|
|
|
|
|
redistribute,
|
|
|
|
|
shard_colwise,
|
|
|
|
|
shard_rowwise,
|
|
|
|
|
sharded_tensor_to_param,
|
|
|
|
|
to_global,
|
2023-06-23 08:07:09 +00:00
|
|
|
to_global_for_customized_distributed_tensor,
|
2023-06-22 03:42:11 +00:00
|
|
|
)
|
|
|
|
|
from .layout import Layout
|
|
|
|
|
from .sharding_spec import ShardingSpec
|
|
|
|
|
|
|
|
|
|
__all__ = [
|
2023-09-19 06:20:26 +00:00
|
|
|
"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",
|
2023-06-22 03:42:11 +00:00
|
|
|
]
|
