mojolearn.MoEBlock

class mojolearn.MoEBlock(hidden_size, intermediate_size, num_experts=8, top_k=2, norm_topk_prob=True, hidden_act='silu', random_state=0, numeric_mode=None)

Bases: object

Sparse mixture-of-experts feed-forward block, as Mixtral’s MixtralSparseMoeBlock.

A router scores every token against num_experts experts; the top_k highest softmax probabilities are kept (renormalized to sum to one when norm_topk_prob), each chosen expert computes down_proj(silu(gate) * up), and the output is the weighted sum. Routing ties go to the lower expert index. Weights keep the Hugging Face layout. Forward only.

The weight arrays are stored as read-only float32 copies: to change a weight, assign a new array (block.router = new) rather than editing it in place.

Not supported: the backward pass, router jitter noise, activations other than SiLU, the auxiliary load-balancing loss, and float64 input.

Parameters:

• hidden_size (int) – Model dimension D.

• intermediate_size (int) – Expert hidden dimension F.

• num_experts (int, default 8) – Number of experts E.

• top_k (int, default 2) – Experts per token; 1 <= top_k <= num_experts, otherwise ValueError.

• norm_topk_prob (bool, default True) – Renormalize the kept routing probabilities.

• hidden_act (str, default ‘silu’) – Expert activation; anything else raises NotImplementedError.

• random_state (int, default 0) – Seed of the N(0, 0.02**2) weight initialization.

• numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Variables:

• router (ndarray, shape (E, D)) – Router weights.

• gate_up_proj (ndarray, shape (E, 2F, D)) – Gate and up projections of every expert.

• down_proj (ndarray, shape (E, D, F)) – Down projection of every expert.

• router_logits_ (ndarray, shape (n_tokens, E)) – Router logits of the last forward.

• selected_experts_ (ndarray of int64, shape (n_tokens, top_k)) – Experts chosen for each token by the last forward.

• routing_weights_ (ndarray, shape (n_tokens, top_k)) – Routing weights of the chosen experts.

Examples

>>> import numpy as np, mojolearn
>>> block = mojolearn.MoEBlock(hidden_size=64, intermediate_size=128, num_experts=8, top_k=2)
>>> x = np.random.default_rng(0).random((4, 16, 64), dtype=np.float32)
>>> y = block(x)

__init__(hidden_size, intermediate_size, num_experts=8, top_k=2, norm_topk_prob=True, hidden_act='silu', random_state=0, numeric_mode=None)

Methods

forward(x)	Apply the block to x of shape (..., hidden_size); also callable as block(x).
load_state_dict(sd)	Load router, gate_up_proj and down_proj from a dict of arrays; returns self.
to_device(a)	Upload a float32 array once so forward keeps its input and outputs on the device.

Attributes

router_logits_	Router logits of the last forward, (n_tokens, num_experts).
routing_weights_	Routing weights of the chosen experts, (n_tokens, top_k).
selected_experts_	Experts chosen for each token by the last forward, int64 (n_tokens, top_k).

to_device(a)

Upload a float32 array once so forward keeps its input and outputs on the device.

When this install has no device path for the block, the array itself is returned. Non-float32 input raises TypeError on that path.

Parameters:

a (ndarray of float32) – Array to upload.

Returns:

device tensor or ndarray

property router_logits_

Router logits of the last forward, (n_tokens, num_experts).

property selected_experts_

Experts chosen for each token by the last forward, int64 (n_tokens, top_k).

property routing_weights_

Routing weights of the chosen experts, (n_tokens, top_k).

load_state_dict(sd)

Load router, gate_up_proj and down_proj from a dict of arrays; returns self.

Shapes must match the block’s, otherwise ValueError.

forward(x)

Apply the block to x of shape (…, hidden_size); also callable as block(x).

Sets router_logits_, selected_experts_ and routing_weights_. float64 input raises TypeError. An input from to_device gives an output on the device.

Parameters:

x (ndarray of float32) – Token activations.

Returns:

y (ndarray of float32) – Output with the shape of x.