1 / 1

Model Ensemble with in_axes

JAX

Easy

Evaluate an ensemble of M small two-layer MLPs on one shared batch of inputs, where the weights of all members are stacked along a leading member axis, and return every member's predictions together with their average.

Member m computes

f_m(x) = \operatorname{relu}\left(x W_1^{(m)} + b_1^{(m)}\right) W_2^{(m)} + b_2^{(m)} ,

where relu(z) = max(z, 0) is applied elementwise. The dict params has the keys w1 of shape (M, D, K), b1 of shape (M, K), w2 of shape (M, K, O) and b2 of shape (M, O), so $W_1^{(m)}$ is params["w1"][m], $b_1^{(m)}$ is params["b1"][m], and so on. The input x has shape (N, D) and holds N rows that every member sees. Here D is the input width, K is the hidden width and O is the output width.

Return a dict with two float32 JAX arrays, outputs of shape (M, N, O), where outputs[m] is $f_m(x)$ , and mean of shape (N, O), the average of outputs over the M members.

Examples

Input: params = {"w1":[[[1,0],[0,1]],[[1,1],[0,1]],[[-1,0],[0,2]]],"b1":[[0,0],[0,-1],[1,0]],"w2":[[[1],[1]],[[2],[-1]],[[0.5],[1]]],"b2":[[0],[1],[-1]]}, x = [[1,2],[-1,0.5]]

Output: {"outputs":[[[3.0],[0.5]],[[1.0],[1.0]],[[3.0],[1.0]]],"mean":[[2.333333],[0.833333]]}

Explanation: Member 0 maps the input (1, 2) to the hidden values relu(1, 2) = (1, 2) and then to 1 + 2 = 3. The three members predict 3, 1 and 3 for that input, so its mean is 7 / 3 = 2.333333.

Input: params = {"w1":[[[1,-1,0],[0,1,1],[1,0,-1]],[[0.5,0.5,0],[-1,0,1],[0,2,0]]],"b1":[[0,0.5,0],[0.1,0,-0.2]],"w2":[[[1,0],[0,1],[1,-1]],[[2,0],[0,-1],[1,1]]],"b2":[[0,0],[0.5,-0.5]]}, x = [[1,-2,0.5]]

Output: {"outputs":[[[1.5,0.0]],[[5.7,-2.0]]],"mean":[[3.6,-1.0]]}

Input: params = {"w1":[[[2,-1]]],"b1":[[-1,1]],"w2":[[[1],[3]]],"b2":[[0.25]]}, x = [[0],[1],[-1]]

Output: {"outputs":[[[3.25],[1.25],[6.25]]],"mean":[[3.25],[1.25],[6.25]]}

Hint 1

jax.vmap(f, in_axes=(0, None)) maps the first argument and shares the second.

Hint 2

in_axes=0 on a dict argument maps every array in the dict.

Hint 3

jnp.mean(outs, axis=0) averages over the ensemble axis.

Requirements

Write the MLP for the weights of a single member and evaluate all members with a function built by jax.vmap whose in_axes maps the stacked weights over axis 0 and marks x as shared with None.
The function must also give the same result when it is wrapped in jax.jit.
Leave params unchanged, with the same keys and arrays.

Constraints

M is between 1 and 5, N is between 1 and 4, D is between 1 and 4, K is between 2 and 6, and O is between 1 and 3.
params has exactly the keys w1, b1, w2 and b2.
Every array in params and x is float32, and all values are finite.

Try Similar Problems

Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium Functional Array UpdateEasy PRNG Key SplittingEasy Masked Mean Without Boolean IndexingEasy Value and Grad of a LossEasy Static Arguments in jitEasy SGD with Momentum over a PytreeEasy Rotary Position EmbeddingEasy KV Cache Update at a PositionEasy Device Mesh and Named ShardingEasy Pad and Mask to Bucket LengthMedium