AllToShardedLinear

Distributed linear layer where result is sharded across the group.

Abstract Signature:

AllToShardedLinear(in_features: int, out_features: int, bias: bool = True, group)

PyTorchApple MLXFlax NNX

PyTorch

API: —

Strategy: Custom / Partial

Apple MLX

API: mlx.nn.AllToShardedLinear

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: —

Strategy: Custom / Partial