gensbi.sharding

Multi-device sharding helpers for GenSBI.

Models annotate their variables with logical axis names through axes() (or partitioned() for explicit initialisers). Annotations are inert: nothing is placed on a device until shard_model() maps the logical names onto the axes of a jax.sharding.Mesh according to a named rule ("replicated", "fsdp" or "tensor"). With no mesh, nothing changes.

Cross-repo contract

LOGICAL_AXES, RULES, MODEL_AXIS and the semantics of axes, partitioned, make_mesh, batch_axes, batch_size_divisor, shard_model, shard_batch, check_batch_divisible and sharding_of are identical to heal_swin_nnx/sharding.py. Neither package imports the other; a test pins the vocabulary against the installed heal_swin_nnx. Change both files together.

Attributes

LOGICAL_AXES
MODEL_AXIS
RULES

Functions

_is_array(x)
axes(*names)	Variable metadata annotating each array dimension with a logical axis.
batch_axes(mesh)	Mesh axis names that carry the batch (every axis except "model").
batch_size_divisor(mesh)	Number of data-parallel shards a batch is split into.
check_batch_divisible(batch, mesh, what)	Raise ValueError if the batch cannot be split evenly across the mesh.
make_mesh([devices, model_axis])	Build a ("data", "model") mesh with Auto axis types.
mesh_context(mesh)	Make mesh the active mesh for the body of a with block.
pad_for_mesh(n, mesh)	Round n up to a multiple of the data-parallel size.
partitioned(init, names)	Wrap an initialiser so its variable carries logical axis names.
replicate(tree, mesh)	Replicate every array leaf of tree over every device of mesh.
resolve_mesh(spec)	Turn a user-facing mesh spec into a mesh.
shard_batch(batch, mesh)	Shard the leading axis of every array leaf of batch over the batch axes.
shard_model(node, mesh[, rule])	Place every variable of node on mesh according to rule.
shard_sample_init(x_init, mesh)	Shard a sampler's initial state (array or pytree) over the batch axes.
sharding_of(node)	Map "a/b/c" state paths of node to the sharding of each leaf.

Module Contents

gensbi.sharding._is_array(x)

Return type:

bool

gensbi.sharding.axes(*names)

Variable metadata annotating each array dimension with a logical axis.

Parameters:

*names (str or None) – One logical axis name (see LOGICAL_AXES) or None per array dimension.

Returns:

Keyword arguments for an nnx.Variable constructor, or the value of a *_metadata argument of an NNX layer. eager_sharding=False keeps construction independent of any mesh.

Return type:

dict

gensbi.sharding.batch_axes(mesh)

Mesh axis names that carry the batch (every axis except "model").

Return type:

tuple[str, …]

gensbi.sharding.batch_size_divisor(mesh)

Number of data-parallel shards a batch is split into.

Return type:

int

gensbi.sharding.check_batch_divisible(batch, mesh, what)

Raise ValueError if the batch cannot be split evenly across the mesh.

Parameters:

• batch (pytree) – Batch whose array leaves share the same leading size.

• mesh (jax.sharding.Mesh)

• what (str) – Label used in the error message (e.g. "train").

Return type:

None

gensbi.sharding.make_mesh(devices=None, model_axis=1)

Build a ("data", "model") mesh with Auto axis types.

Parameters:

• devices (sequence of jax.Device, optional) – Devices to use; defaults to jax.devices().

• model_axis (int, optional) – Size of the "model" axis; the "data" axis gets len(devices) // model_axis.

Return type:

jax.sharding.Mesh

Raises:

ValueError – If model_axis does not divide the device count.

gensbi.sharding.mesh_context(mesh)

Make mesh the active mesh for the body of a with block.

jax.set_mesh(mesh) is entered when the with block is entered (not when this function is called) and left on exit. With mesh=None this is a no-op context.

Parameters:

mesh (jax.sharding.Mesh or None) – Mesh to activate.

gensbi.sharding.pad_for_mesh(n, mesh)

Round n up to a multiple of the data-parallel size.

Returns:

(padded, n).

Return type:

tuple of int

Parameters:

n (int)

gensbi.sharding.partitioned(init, names)

Wrap an initialiser so its variable carries logical axis names.

Parameters:

• init (callable) – Initialiser to wrap.

• names (tuple of (str or None)) – Logical axis name per array dimension.

Returns:

nnx.with_partitioning(init, names, eager_sharding=False).

Return type:

callable

gensbi.sharding.replicate(tree, mesh)

Replicate every array leaf of tree over every device of mesh.

Used for inputs that must live on the mesh without being split, such as a sampler’s condition and node ids. Arrays committed to a single device are moved; non-array leaves are returned untouched.

Parameters:

• tree (pytree) – Arrays (jax.Array or numpy.ndarray) and arbitrary other leaves.

• mesh (jax.sharding.Mesh) – Target mesh.

Returns:

Same structure, array leaves placed with NamedSharding(mesh, P()).

Return type:

pytree

gensbi.sharding.resolve_mesh(spec)

Turn a user-facing mesh spec into a mesh.

None -> None; a Mesh -> itself; "auto" -> make_mesh(); a dict -> make_mesh(**spec). Anything else raises ValueError.

Return type:

jax.sharding.Mesh | None

gensbi.sharding.shard_batch(batch, mesh)

Shard the leading axis of every array leaf of batch over the batch axes.

Non-array leaves are returned untouched.

gensbi.sharding.shard_model(node, mesh, rule='replicated')

Place every variable of node on mesh according to rule.

Works in place and is idempotent. node is any NNX graph node (modules, nnx.Optimizer, EMA containers). Variables without annotation are replicated; variables with an unknown logical name, a dimension not divisible by the mapped mesh axis, or a mapped axis missing from mesh fall back to replicated and are reported in a single warning. "replicated" ignores annotations (so never warns about them); under other rules unannotated nnx.Param s are replicated and listed in that warning too.

Parameters:

• node (nnx graph node) – Object whose state is placed.

• mesh (jax.sharding.Mesh) – Target mesh.

• rule ({"replicated", "fsdp", "tensor"}, optional) – Logical-to-mesh axis rule. Anything but "replicated" is experimental and warns on every call.

Raises:

ValueError – If rule is not a key of RULES.

Return type:

None

gensbi.sharding.shard_sample_init(x_init, mesh)

Shard a sampler’s initial state (array or pytree) over the batch axes.

gensbi.sharding.sharding_of(node)

Map "a/b/c" state paths of node to the sharding of each leaf.

Return type:

dict[str, jax.sharding.Sharding]

gensbi.sharding.LOGICAL_AXES: tuple[str, ...] = ('embed', 'mlp', 'heads', 'vocab', 'in', 'out')

gensbi.sharding.MODEL_AXIS = 'model'

gensbi.sharding.RULES: dict[str, tuple[tuple[str, str | None], ...]]