Skip to content

← Flax NNX step 88 of 100

Medium Primitives

NNX Parameter Counts and Ideal Shares

Jump to the task

This exercise builds an NNX Linear layer, counts its scalar parameters, and divides the count by a hypothetical number of shards. It does not save a checkpoint, create device shards or call Orbax.

Count the kernel and bias

For input width D and output width F, default nnx.Linear(D, F) has a kernel of shape (D, F) and a bias of shape (F,). Its parameter count is therefore D * F + F, or (D + 1) * F.

model = nnx.Linear(D, F, rngs=nnx.Rngs(seed))
total = model.kernel[...].size + model.bias[...].size

The [...] access reads the underlying parameter arrays. There are two array parameters, but usually many more scalar values. Batch size and input values do not affect the count. The seed changes initial kernel values, not its shape.

Divide without claiming a physical layout

With D = 4 and F = 4, the total is 20 values. Dividing by two gives an ideal average of 10 values per shard. With six values and four shards, the average is 1.5. No physical shard contains half a scalar: the quotient is accounting, not a partitioning plan.

Actual array partitioning depends on dimension sizes and placement rules. Some parameters can be replicated. The quotient cannot determine file count, allocation per host, padding requirements or checkpoint speed.

Counts are not bytes. Twenty float32 values have an 80-byte raw payload, but checkpoint size can also reflect metadata, other state, storage dtype and compression. Orbax storage chunks and files need not match device shards one-to-one. More writers do not guarantee proportional speedup, and splitting data is not redundancy against losing a shard.

Problem

Implement orbax_save_count(seed, x, features, num_shards):

  1. Convert seed, features and shard count to integers.
  2. Create default nnx.Linear(x.shape[-1], features, rngs=nnx.Rngs(seed)).
  3. Count scalar elements in its kernel and bias, retaining the default bias.
  4. Divide their sum by the shard count using /, without rounding.
  5. Return a floating-point JAX array of shape (2,) containing [total_params, ideal_share].

x is a finite float32 JAX matrix of shape (B, D), with positive B and D. Width and shard count are positive integer-valued scalars; seed is nonnegative and integer-valued. These scalars may arrive as floats. Leave x unchanged. Do not multiply by bytes per value or write files.

The result tests counts, not model construction or checkpointing. A formula can produce the same output. The diagram below illustrates a separate example of array placement; this function does not construct that mesh.

Loading visualization…