FeedForward

Feedforward layer with activation (usually Linear -> Act -> Linear).

Abstract Signature:

FeedForward(in_features: int, out_features: int, hidden_dims: int, activation: str = relu)

PyTorchKerasFlax NNXPaxML / Praxis

PyTorch

API: —

Strategy: Custom / Partial

Keras

API: —

Strategy: Custom / Partial

Flax NNX

API: —

Strategy: Custom / Partial

PaxML / Praxis

API: paxml.layers.FeedForward

Strategy: Direct Mapping

Official Docs ↗