WIP

Binyang2014 · Binyang2014 · commit afb403ee1c98 · 2024-12-27T23:53:55.000Z
diff --git a/msccl/collectives.py b/msccl/collectives.py
diff --git a/msccl/language/collectives.py b/msccl/language/collectives.py
@@ -1,4 +1,6 @@
-from msccl.language import *
+
+from msccl.language.buffer import Buffer
+from msccl.language.chunk import Chunk, ReduceChunk
 
 
 class Collective:
diff --git a/msccl/language/mscclpp/__init__.py b/msccl/language/mscclpp/__init__.py
@@ -2,7 +2,7 @@
 # Licensed under the MIT License.
 
 from dataclasses import dataclass
-from msccl.collectives import Collective
+from msccl.language.collectives import Collective
 from msccl.language.buffer import *
 from msccl.language.types import ChannelType, ChunkRef, ReplicationPolicy, Threadblock
 from msccl.language.mscclpp.ir import *
