Skip to content

← Flax step 91 of 100

Medium Primitives

nn.with_partitioning Annotation

Jump to the task

Attach axis metadata to a parameter

Create a Dense layer whose kernel initializer is wrapped with nn.with_partitioning. Initialize the layer and return the kernel's full shape, [input_width, output_width].

The wrapper stores partitioning metadata alongside an array. It does not change the Dense formula, x @ kernel + bias, or the kernel's logical shape. This exercise uses no device mesh and does not distribute the model across multiple devices.

Wrap the initializer

Use this initializer for the kernel:

nn.with_partitioning(nn.initializers.lecun_normal(), (None, "model"))

It calls the underlying initializer and returns a nn.Partitioned value containing the array and its axis names. For a kernel of shape (D, F), (None, "model") means that the first array dimension is not partitioned and the output dimension is associated with the mesh axis named model when a suitable mesh is used.

This does not mean the entire array is replicated: the output dimension may be split. Mesh axes that are not used by the specification replicate the resulting blocks. Without a mesh, these names remain metadata.

Keep the Dense bias enabled with its ordinary zero initializer. The task annotates only the kernel, so the stored bias is a plain array of shape (F,).

Inspect the box and the array separately

The value in the variable tree is a metadata object, not a general drop-in array. Read its .names and .value, or use the appropriate unboxing utilities. Linen normally unboxes parameters when the layer uses them in arithmetic.

This example shows the stored kernel and confirms that the metadata does not change its numerical value in the no-mesh setup:

import jax
import jax.numpy as jnp
import flax.linen as nn
from jax.sharding import PartitionSpec as P

x = jnp.ones((2, 4))
key = jax.random.PRNGKey(0)
annotated = nn.Dense(
    features=8,
    kernel_init=nn.with_partitioning(
        nn.initializers.lecun_normal(), (None, "model")
    ),
)
variables = annotated.init(key, x)
kernel = variables["params"]["kernel"]
assert isinstance(kernel, nn.Partitioned)
assert kernel.names == (None, "model")
assert kernel.value.shape == (4, 8)
assert variables["params"]["bias"].shape == (8,)
assert nn.get_partition_spec(variables)["params"]["kernel"] == P(None, "model")

plain = nn.Dense(features=8)
plain_variables = plain.init(key, x)
assert jnp.allclose(kernel.value, plain_variables["params"]["kernel"])
assert jnp.allclose(annotated.apply(variables, x), plain.apply(plain_variables, x))

Metadata needs a sharding context

nn.get_partition_spec extracts a matching tree of PartitionSpec objects. A specification names axes but does not identify physical devices. For explicit placement, combine a specification with a device mesh using jax.sharding.NamedSharding before passing it to jax.device_put.

With an active or explicitly supplied mesh, unboxing a Partitioned value can apply a sharding constraint. Mesh names, sizes and array dimensions must then be compatible. Annotation alone does not guarantee a particular layout in every downstream tool, or force an unannotated parameter to be replicated in every compilation.

Do not interpret the wrapper as a promise to allocate a full weight matrix on one device before sharding it. Physical initialization and placement depend on the surrounding JAX execution and sharding setup.

Problem

Implement partitioned_dense_kernel_shape(seed, x, features):

  1. Convert features to Python integer F and create jax.random.PRNGKey(int(seed)).
  2. Construct nn.Dense(features=F) with the wrapped kernel initializer shown above. Keep the default bias initializer.
  3. Initialize the layer directly with the key and x.
  4. Return jnp.array([float(x.shape[-1]), float(F)]).

x is a two-dimensional float array of shape (N, D). N, D and F are positive. The output is a one-dimensional float array of shape (2,), containing [D, F], not the Dense activations or a per-device shard shape. Batch size N does not enter the kernel shape.

Keep seed and features static if JIT-compiling the entrypoint. The returned dimensions alone cannot prove that the annotation was created; the initialization and inspection above are part of what this exercise asks you to learn.

The diagram illustrates how (None, "model") would divide an (8, 12) kernel on a hypothetical data=2, model=4 mesh: blocks have shape (8, 3), with two replicas of each block along data. It does not describe devices allocated by this function.

Loading visualization…