chickadee » nanograd-components » make-am-mlp-mixer

make-am-mlp-mixer seq-len hidden-dim n-blocks #!key (token-expansion 4) (channel-expansion 4) (use-rmsnorm #t) (norm-eps 1e-06) (num-classes #f) (dtype 'f32) (name AM-MLPMixer)procedure

Creates the lazy model. Without num-classes, the output is the [T*B, D] feature matrix, as in the strict model. With num-classes, the model averages the features over the tokens, and a dense head gives [B, num-classes] logits. The head's parameters come last.