Add distributed layers to nn top-level

ml-explore · Jul 15, 2024 · 071fb86 · 071fb86
1 parent 84e1444
commit 071fb86
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/mlx/nn/layers/__init__.py b/python/mlx/nn/layers/__init__.py
@@ -55,6 +55,12 @@
 from mlx.nn.layers.base import Module
 from mlx.nn.layers.containers import Sequential
 from mlx.nn.layers.convolution import Conv1d, Conv2d, Conv3d
+from mlx.nn.layers.distributed import (
+    AllToShardedLinear,
+    QuantizedAllToShardedLinear,
+    QuantizedShardedToAllLinear,
+    ShardedToAllLinear,
+)
 from mlx.nn.layers.dropout import Dropout, Dropout2d, Dropout3d
 from mlx.nn.layers.embedding import Embedding
 from mlx.nn.layers.linear import Bilinear, Identity, Linear