StaticArray migration guide (v0.2 → v0.3.0)#
Changed in version v0.2: StaticArray introduced as the
typed wrapper for non-state arrays carried via
SimulationNode.static_data.
Bare arrays (anything with shape/dtype but not wrapped) were
coerced to StaticArray(value=arr) with a one-time
FutureWarning per (node × key) pair.
Changed in version v0.3.0: The bare-array coercion path was removed. Putting a bare
ndarray / jnp.ndarray into static_data now raises
MigrationError immediately.
Added in version v0.3.0: New replication="partition" variant for graph-partitioned
sharding (see Sharding topology: structured vs unstructured).
Why this changed#
StaticArray is the declaration site for how an array gets
materialised under sharding:
Replication |
What the wrapper does |
|---|---|
|
Every device gets the full array. |
|
Per-device slice along |
|
Per-device slice via a partition assignment (graph — paired with |
A bare array carries no replication policy, so the v0.2 coercion
path silently chose "replicate". That was harmless on a single
device but a bug magnet under sharding — the user couldn’t see
why their stencil mask was bigger than expected per device.
v0.3.0 makes the declaration mandatory.
What to do#
Wrap your static_data values explicitly#
class MyNode(SimulationNode):
@property
def static_data(self) -> dict:
return {
"weights": self._weights_array, # bare array
"mask": self._mask_array, # bare array
}
from maddening.core.static_data import StaticArray
class MyNode(SimulationNode):
@property
def static_data(self) -> dict:
return {
# Lookup table that every shard needs in full.
"weights": StaticArray(self._weights_array),
# Per-shard slab along axis 0 — paired with ShardedStencilNode
# whose axis_map sends a mesh axis to the same spatial axis 0.
"mask": StaticArray(
self._mask_array,
replication="shard", shard_axis=0,
),
}
Scalars, strings, and tuples pass through unchanged#
Only array-like values (with shape AND dtype) trigger the
migration error. Plain Python scalars, strings, and tuples stay
bare:
@property
def static_data(self) -> dict:
return {
"weights": StaticArray(self._w), # wrapped
"n_iters": 42, # scalar — no wrap needed
"label": "demo", # string — no wrap needed
"shape": (8, 8, 8), # tuple — no wrap needed
}
Nested structures stay unsupported#
A list/dict of arrays under a single static_data key still raises
at StaticArray.__post_init__. Unfold into multiple top-level keys:
# WRONG -- raises TypeError
"masks": [StaticArray(m1), StaticArray(m2)]
# RIGHT
"mask_a": StaticArray(m1),
"mask_b": StaticArray(m2),
The new "partition" variant for unstructured sharding#
If you’re authoring an unstructured / graph-partitioned node for
ShardedUnstructuredNode,
declare any per-cell static data with the new "partition"
replication and a partition-assignment array:
import numpy as np
from maddening.core.static_data import StaticArray
class FVMNode(SimulationNode):
def __init__(self, *, cell_volumes, partition_assignment, ...):
self._cell_volumes = cell_volumes # shape (n_global,)
self._pa = partition_assignment.astype(np.int32) # shape (n_global,)
...
@property
def static_data(self) -> dict:
return {
"cell_volumes": StaticArray(
self._cell_volumes,
replication="partition",
partition_assignment=self._pa,
),
}
By convention the partitioned axis is axis 0 of value (the
global-cell axis). The wrapper validates the assignment’s length
matches.
See Sharding topology: structured vs unstructured for the full unstructured-sharding contract.
Auto-migration tooling#
The MigrationError raised from
coerce_static_data_value carries structured detail:
from maddening.warnings import MigrationError
try:
my_node.static_data_hash()
except MigrationError as err:
print(f"api_name: {err.api_name}")
print(f"replacement: {err.replacement}")
print(f"guide: {err.migration_guide}")
The api_name follows the convention bare-array-in-static_data (<node>.<key>), so a tool can grep MIME / MICROROBOTICA for the
specific offending entries.