LinearGeneral

Flexible linear layer with axis specification.

Abstract Signature:

LinearGeneral(in_features: int, out_features: int, axis: Union[int, Tuple[int]])

PyTorchKerasFlax NNXPaxML / Praxis

PyTorch

API: —

Strategy: Plugin (linear_general_adapter)

Keras

API: keras.layers.Dense

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.LinearGeneral

Strategy: Direct Mapping

Official Docs ↗

PaxML / Praxis

API: praxis.layers.Linear

Strategy: Direct Mapping

Official Docs ↗