Skip to content

Commit

Permalink
feat: add barrier to bagua.torch_api module (#698)
Browse files Browse the repository at this point in the history
  • Loading branch information
wangraying authored Apr 15, 2023
1 parent 629fd71 commit b70fcf2
Showing 1 changed file with 13 additions and 14 deletions.
27 changes: 13 additions & 14 deletions bagua/torch_api/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,7 @@
"""
The Bagua communication library PyTorch interface.
"""
from distutils.errors import (
DistutilsPlatformError,
)
from distutils.errors import DistutilsPlatformError

try:
import torch
Expand All @@ -29,32 +27,33 @@
init_process_group,
send,
recv,
broadcast,
reduce,
reduce_inplace,
gather,
gather_inplace,
scatter,
scatter_inplace,
allreduce,
allreduce_inplace,
allgather,
allgather_inplace,
allreduce,
allreduce_inplace,
alltoall,
alltoall_inplace,
alltoall_v,
alltoall_v_inplace,
barrier,
broadcast,
gather,
gather_inplace,
reduce,
reduce_inplace,
reduce_scatter,
reduce_scatter_inplace,
scatter,
scatter_inplace,
ReduceOp,
)
from .distributed import BaguaModule # noqa: F401
from .tensor import BaguaTensor # noqa: F401
from .env import ( # noqa: F401
get_rank,
get_world_size,
get_local_rank,
get_local_size,
get_rank,
get_world_size,
)
from . import contrib # noqa: F401
from . import communication # noqa: F401
Expand Down

0 comments on commit b70fcf2

Please sign in to comment.