chickadee » nanograd-components » make-mlp-mixer

make-mlp-mixer seq-len hidden-dim n-blocks #!key (token-expansion 4) (channel-expansion 4) (use-rmsnorm #t) (dtype 'f32) nameprocedure

Creates a model of n-blocks blocks for one (seq-len, hidden-dim) sample. The output has the same shape as the input. The token MLP has (* token-expansion seq-len) hidden units, and the channel MLP (* channel-expansion hidden-dim). With use-rmsnorm: #f, the blocks use LayerNorm.