ImplementClaim probe

Switch: auxiliary load-balancing loss

Implement `run(fractions, probabilities, alpha)` returning the auxiliary load-balancing loss for a Switch/GShard-style router: alpha * N * sum(f_i * P_i), where N is the number of experts, f_i is the fraction of tokens routed to expert i, P_i is the fraction of router dispatch probability assigned to expert i, and alpha is the loss coefficient. Return the float value.

python12 mindifficulty 3/4

From Switch Transformers

Sample tests

run([[0.25,0.25,0.25,0.25],[0.25,0.25,0.25,0.25],0.01])→0.01
run([[0.5,0.3,0.15,0.05],[0.4,0.3,0.2,0.1],0.01])→0.013

Hints

The auxiliary loss pushes the router toward uniform load so experts do not starve or overflow.
L_aux = alpha * N * sum over i of (f_i * P_i).
Perfect balance gives f_i = P_i = 1/N for every expert, making the loss exactly alpha.
fractions and probabilities should each sum to 1 (they are normalized distributions).

Hints

Code runs in an isolated sandbox on the OpenSandbox judge. No model key required.