zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-12 12:16:43 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add distributed layers to nn top-level

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-07-15 13:23:51 -07:00
parent 060e1c9f92
commit a8b3da7946
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
python/mlx/nn/layers/__init__.py
View File

@ -60,6 +60,12 @@ from mlx.nn.layers.convolution_transpose import (
ConvTranspose2d, ConvTranspose2d,
ConvTranspose3d, ConvTranspose3d,
) )
from mlx.nn.layers.distributed import (
AllToShardedLinear,
QuantizedAllToShardedLinear,
QuantizedShardedToAllLinear,
ShardedToAllLinear,
)
from mlx.nn.layers.dropout import Dropout, Dropout2d, Dropout3d from mlx.nn.layers.dropout import Dropout, Dropout2d, Dropout3d
from mlx.nn.layers.embedding import Embedding from mlx.nn.layers.embedding import Embedding
from mlx.nn.layers.linear import Bilinear, Identity, Linear from mlx.nn.layers.linear import Bilinear, Identity, Linear