GShardSharedEmbeddingSoftmax

Softmax layer with embedding lookup and Gaussian init used in GShard.

Abstract Signature:

GShardSharedEmbeddingSoftmax(in_features: int, num_classes: int)

PyTorchKerasFlax NNXPaxML / Praxis

PyTorch

API: —

Strategy: Custom / Partial

Keras

API: —

Strategy: Custom / Partial

Flax NNX

API: —

Strategy: Custom / Partial

PaxML / Praxis

API: paxml.layers.GShardSharedEmbeddingSoftmax

Strategy: Direct Mapping

Official Docs ↗