Evaluate MSE in Equal Chunks

Problem

Implement sharded_eval_loss(seed, x, y, num_shards):

1. Convert seed and num_shards to Python integers.
2. Use the TinyDense module below. Initialize it once with jax.random.PRNGKey(seed) and x[:1], without splitting the key.
3. Split x and y into num_shards equal chunks along axis 0. Keep each input row paired with its target.
4. Apply the model to each input chunk. Flatten the (chunk_size, 1) predictions to (chunk_size,) before subtracting the targets.
5. Compute each chunk's mean squared error, then average those losses.
6. Return a JAX array of shape (1,) containing the average loss.

import flax.linen as nn

class TinyDense(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Dense(1)(x)

Keep this wrapper and the unnamed Dense child: its parameter path is Dense_0, which affects Flax's random initialization. A standalone nn.Dense(1) initialized with the same seed need not get the same kernel. Use the default kernel initializer and zero bias.

Inputs: x has shape (N, D) and y has shape (N,), with positive N and D. The positive chunk count divides N exactly. The seed and chunk count are integer-valued scalars that may arrive as floats. Output: [average_loss], a one-dimensional array with one element.

Why equal chunks matter

Suppose the prediction errors are [1, 3, 5, 7]. The squared errors are [1, 9, 25, 49]. Two equal chunks have losses 5 and 37, and their average is 21, the same as the mean of all four squared errors.

For unequal chunks, weight by the number of examples. With errors [0, 0, 6], chunks of sizes two and one have losses 0 and 36. Their unweighted average is 18, but the correct MSE is 36 / 3 = 12. Summing squared errors and dividing by the total example count handles unequal chunk sizes. This task only requires equal, nonempty chunks; do not discard examples to make the split work.

Avoid accidental broadcasting

Subtracting targets of shape (C,) from predictions of shape (C, 1) produces a (C, C) array. That compares every prediction with every target. Reshape predictions to (C,) first so each row contributes one squared error.

Chunking, compilation, and memory

For this Dense layer, the input feature width determines the parameter shapes. A single row is enough for initialization. Reusing the variables tree lets each chunk use the same kernel and bias.

Chunking can reduce the size of intermediate activations. It does not by itself release the full x and y arrays passed to this function. JAX also dispatches work asynchronously, so a Python loop alone does not guarantee that each device computation finishes before the next starts or that memory is freed between iterations.

A fixed Python loop inside jax.jit is traced into the compiled function; it does not necessarily compile a separate program for every iteration. vmap can express a batched calculation too, but neither approach has a universal speed or memory advantage. Measure your actual workload.

Keep seed and num_shards static if you JIT this function. Keep the final loss as a JAX value, for example with jnp.reshape(avg, (1,)). Converting it to a Python float would prevent tracing that operation under JIT and break differentiation through the returned loss.

The diagram illustrates two equal row chunks, each (4, 12). It is a layout analogy; the graded function does not allocate a device mesh.