FFN

ffn

Feed-forward / MLP blocks (vanilla + gated variants).

FeedForward

FeedForward(dim: int, dim_out: int | None = None, mult: float = 4.0, *, inner_dim: int | None = None, activation: str | Callable[[array], array] = 'gelu', bias: bool = True)

Bases: Module

Classic 2-Linear feed-forward block.

Parameters:

Name	Type	Description	Default
dim	int	Input dimension.	required
dim_out	int | None	Output dimension. Defaults to dim (square FFN).	None
mult	float	Multiplier on dim to derive the inner dimension when inner_dim is not supplied. Defaults to 4.0.	4.0
inner_dim	int | None	Explicit inner dimension. Overrides mult.	None
activation	str | Callable[[array], array]	Either a name ("gelu", "gelu_approx", "silu", "relu") or a callable. Defaults to "gelu".	'gelu'
bias	bool	Whether the two Linears include bias. Defaults to True.	True

Weight keys

proj_in.weight, proj_in.bias (if bias=True), proj_out.weight, proj_out.bias (if bias=True).

GatedFFN

GatedFFN(dim: int, dim_out: int | None = None, mult: float = 4.0, *, inner_dim: int | None = None, gate_activation: str | Callable[[array], array] = 'gelu', bias: bool = False)

Bases: Module

Gated 3-Linear feed-forward block (GeGLU / SwiGLU).

Computes down_proj( activation(gate_proj(x)) * up_proj(x) ).

Parameters:

Name	Type	Description	Default
dim	int	Input dimension.	required
dim_out	int | None	Output dimension. Defaults to dim.	None
mult	float	Multiplier on dim for the inner dimension when inner_dim is not supplied. Defaults to 4.0. (Many LLaMA-family models use mult ≈ 2.67 to keep param count comparable to a 4× standard FFN despite the extra projection — caller supplies inner_dim directly in that case.)	4.0
inner_dim	int | None	Explicit inner dimension. Overrides mult.	None
gate_activation	str | Callable[[array], array]	"gelu" for GeGLU (Shazeer 2020), "silu" for SwiGLU (LLaMA / PaLM). Also accepts "gelu_approx" or a callable.	'gelu'
bias	bool	Whether the three Linears include bias. Defaults to False (matches LLaMA / ERNIE).	False

Weight keys

gate_proj.weight, up_proj.weight, down_proj.weight (and corresponding .bias keys if bias=True).

GeGLU

GeGLU(dim: int, dim_out: int | None = None, mult: float = 4.0, *, inner_dim: int | None = None, bias: bool = False) -> GatedFFN

:class:GatedFFN with gate_activation="gelu" (Shazeer 2020).

SwiGLU

SwiGLU(dim: int, dim_out: int | None = None, mult: float = 4.0, *, inner_dim: int | None = None, bias: bool = False) -> GatedFFN

:class:GatedFFN with gate_activation="silu" (LLaMA / PaLM).

ffn

Feed-forward / MLP blocks used by transformer architectures.

Two flavours cover the bulk of real-world ports:

• :class:FeedForward — the classic 2-Linear MLP Linear(d, d_inner) → activation → Linear(d_inner, d_out). Configurable activation (gelu, gelu_approx, silu, relu). Used by LTX-Video, vanilla DiT, T5, most early transformers.

• :class:GatedFFN — the gated 3-Linear MLP down(act(gate(x)) * up(x)). gate_activation selects the variant: "gelu" is GeGLU (Shazeer 2020), "silu" is SwiGLU (LLaMA / PaLM family). Used by ERNIE-Image, modern LLM-style transformers, mixture-of-experts experts.

Weight-key conventions follow the most common public naming:

• :class:FeedForward → proj_in.{weight,bias}, proj_out.{weight,bias} (LTX / vanilla naming).
• :class:GatedFFN → gate_proj.weight, up_proj.weight, down_proj.weight (LLaMA / HF naming).

Ports that load from diffusers-style net.0/net.2 or Megatron-style linear_fc1/linear_fc2 should remap keys at load time (mlx_arsenal.layout.load_safetensors(..., key_map=...)); arsenal's job is to provide the math, not chase every upstream's bikeshed.

FeedForward

FeedForward(dim: int, dim_out: int | None = None, mult: float = 4.0, *, inner_dim: int | None = None, activation: str | Callable[[array], array] = 'gelu', bias: bool = True)

Bases: Module

Classic 2-Linear feed-forward block.

Parameters:

Name	Type	Description	Default
dim	int	Input dimension.	required
dim_out	int | None	Output dimension. Defaults to dim (square FFN).	None
mult	float	Multiplier on dim to derive the inner dimension when inner_dim is not supplied. Defaults to 4.0.	4.0
inner_dim	int | None	Explicit inner dimension. Overrides mult.	None
activation	str | Callable[[array], array]	Either a name ("gelu", "gelu_approx", "silu", "relu") or a callable. Defaults to "gelu".	'gelu'
bias	bool	Whether the two Linears include bias. Defaults to True.	True

Weight keys

proj_in.weight, proj_in.bias (if bias=True), proj_out.weight, proj_out.bias (if bias=True).

GatedFFN

GatedFFN(dim: int, dim_out: int | None = None, mult: float = 4.0, *, inner_dim: int | None = None, gate_activation: str | Callable[[array], array] = 'gelu', bias: bool = False)

Bases: Module

Gated 3-Linear feed-forward block (GeGLU / SwiGLU).

Computes down_proj( activation(gate_proj(x)) * up_proj(x) ).

Parameters:

Name	Type	Description	Default
dim	int	Input dimension.	required
dim_out	int | None	Output dimension. Defaults to dim.	None
mult	float	Multiplier on dim for the inner dimension when inner_dim is not supplied. Defaults to 4.0. (Many LLaMA-family models use mult ≈ 2.67 to keep param count comparable to a 4× standard FFN despite the extra projection — caller supplies inner_dim directly in that case.)	4.0
inner_dim	int | None	Explicit inner dimension. Overrides mult.	None
gate_activation	str | Callable[[array], array]	"gelu" for GeGLU (Shazeer 2020), "silu" for SwiGLU (LLaMA / PaLM). Also accepts "gelu_approx" or a callable.	'gelu'
bias	bool	Whether the three Linears include bias. Defaults to False (matches LLaMA / ERNIE).	False

Weight keys

gate_proj.weight, up_proj.weight, down_proj.weight (and corresponding .bias keys if bias=True).

GeGLU

GeGLU(dim: int, dim_out: int | None = None, mult: float = 4.0, *, inner_dim: int | None = None, bias: bool = False) -> GatedFFN

:class:GatedFFN with gate_activation="gelu" (Shazeer 2020).

SwiGLU

SwiGLU(dim: int, dim_out: int | None = None, mult: float = 4.0, *, inner_dim: int | None = None, bias: bool = False) -> GatedFFN

:class:GatedFFN with gate_activation="silu" (LLaMA / PaLM).