diff --git a/crazyflow/control/mellinger/control.py b/crazyflow/control/mellinger/control.py
index b122970f..9f0ef661 100644
--- a/crazyflow/control/mellinger/control.py
+++ b/crazyflow/control/mellinger/control.py
@@ -21,7 +21,7 @@
 
 from crazyflow.control.core import controllable, load_params
 from crazyflow.control.transform import force2pwm, motor_force2rotor_vel, pwm2force
-from crazyflow.utils import leaf_replace
+from crazyflow.utils import CORE_NDIM_KEY, leaf_replace
 
 if TYPE_CHECKING:
     from jax import Device
@@ -311,19 +311,19 @@ def force_torque2rotor_vel(
 
 @dataclass
 class MellingerStateData:
-    cmd: Array  # (N, M, 13)
+    cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 13)
     """Full state control command for the drone.
 
     A command consists of [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate].
     We currently do not use the acceleration and angle rate components. This is subject to change.
     """
-    staged_cmd: Array  # (N, M, 13)
+    staged_cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 13)
     """Staging buffer to store the most recent command until the next controller tick."""
-    steps: Array  # (N, 1)
+    steps: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, 1)
     """Last simulation steps that the state control command was applied."""
     freq: int = field(pytree_node=False)
     """Frequency of the state control command."""
-    pos_err_i: Array  # (N, M, 3)
+    pos_err_i: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Integral errors of the state control command."""
     # Parameters for the state controller
     params: dict[str, Array]
@@ -344,20 +344,20 @@ def create(
 
 @dataclass
 class MellingerAttitudeData:
-    cmd: Array  # (N, M, 4)
+    cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
     """Full attitude control command for the drone.
 
     A command consists of [roll, pitch, yaw, collective thrust].
     """
-    staged_cmd: Array  # (N, M, 4)
+    staged_cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
     """Staging buffer to store the most recent command until the next controller tick."""
-    steps: Array  # (N, 1)
+    steps: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, 1)
     """Last simulation steps that the attitude control command was applied."""
     freq: int = field(pytree_node=False)
     """Frequency of the attitude control command."""
-    r_int_error: Array  # (N, M, 3)
+    r_int_error: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Integral errors of the attitude control command."""
-    last_ang_vel: Array  # (N, M, 3)
+    last_ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Last angular velocity of the drone."""
     # Parameters for the attitude controller
     params: dict[str, Array]
@@ -384,14 +384,14 @@ def create(
 
 @dataclass
 class MellingerForceTorqueData:
-    cmd: Array  # (N, M, 4)
+    cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
     """Force-torque command for the drone.
 
     A command consists of [fz, tx, ty, tz].
     """
-    staged_cmd: Array  # (N, M, 4)
+    staged_cmd: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
     """Staging buffer to store the most recent command until the next controller tick."""
-    steps: Array  # (N, 1)
+    steps: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, 1)
     """Last simulation steps that the force and torque control command was applied."""
     freq: int = field(pytree_node=False)
     """Frequency of the force and torque control command."""
diff --git a/crazyflow/dynamics/first_principles/dynamics.py b/crazyflow/dynamics/first_principles/dynamics.py
index 70ee0fdc..1295ff10 100644
--- a/crazyflow/dynamics/first_principles/dynamics.py
+++ b/crazyflow/dynamics/first_principles/dynamics.py
@@ -22,13 +22,13 @@
 import jax.numpy as jnp
 from array_api_compat import array_namespace
 from array_api_compat import device as xp_device
-from flax.struct import dataclass
+from flax.struct import dataclass, field
 from scipy.spatial.transform import Rotation as R
 
 import crazyflow.dynamics.symbols as symbols
 from crazyflow.dynamics.core import load_params, supports
 from crazyflow.dynamics.utils import rotation
-from crazyflow.utils import to_xp
+from crazyflow.utils import CORE_NDIM_KEY, to_xp
 
 if TYPE_CHECKING:
     from jax import Device
@@ -298,27 +298,27 @@ def symbolic_dynamics(
 
 @dataclass
 class Params:
-    mass: Array  # (N, M, 1)
+    mass: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 1)
     """Mass of the drone."""
-    L: Array  # (N, M, 1)
+    L: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Arm length of the drone."""
-    prop_inertia: Array  # (N, M, 1)
+    prop_inertia: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Inertia of the propeller."""
-    gravity_vec: Array  # (N, M, 3)
+    gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Gravity vector of the drone."""
-    J: Array  # (N, M, 3, 3)
+    J: Array = field(metadata={CORE_NDIM_KEY: 2})  # (N, M, 3, 3)
     """Inertia matrix of the drone."""
-    J_inv: Array  # (N, M, 3, 3)
+    J_inv: Array = field(metadata={CORE_NDIM_KEY: 2})  # (N, M, 3, 3)
     """Inverse of the inertia matrix of the drone."""
-    rpm2thrust: Array  # (N, M, 1)
+    rpm2thrust: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Force constant of the drone."""
-    rpm2torque: Array  # (N, M, 1)
+    rpm2torque: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Torque constant of the drone."""
-    mixing_matrix: Array  # (N, M, 3, 4)
+    mixing_matrix: Array = field(metadata={CORE_NDIM_KEY: 2})  # (3, 4)
     """Mixing matrix of the drone."""
-    drag_matrix: Array  # (N, M, 3, 3)
+    drag_matrix: Array = field(metadata={CORE_NDIM_KEY: 2})  # (3, 3)
     """Drag matrix of the drone."""
-    rotor_dyn_coef: Array  # (N, M, 4)
+    rotor_dyn_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (4,)
     """Rotor speed dynamics time constant of the drone."""
 
     @staticmethod
diff --git a/crazyflow/dynamics/so_rpy/dynamics.py b/crazyflow/dynamics/so_rpy/dynamics.py
index a5bffb15..4af83291 100644
--- a/crazyflow/dynamics/so_rpy/dynamics.py
+++ b/crazyflow/dynamics/so_rpy/dynamics.py
@@ -21,13 +21,13 @@
 import jax.numpy as jnp
 from array_api_compat import array_namespace
 from array_api_compat import device as xp_device
-from flax.struct import dataclass
+from flax.struct import dataclass, field
 from scipy.spatial.transform import Rotation as R
 
 import crazyflow.dynamics.symbols as symbols
 from crazyflow.dynamics.core import load_params, supports
 from crazyflow.dynamics.utils import rotation
-from crazyflow.utils import to_xp
+from crazyflow.utils import CORE_NDIM_KEY, to_xp
 
 if TYPE_CHECKING:
     from jax import Device
@@ -316,31 +316,31 @@ def symbolic_dynamics_euler(
 
 @dataclass
 class Params:
-    mass: Array  # (N, M, 1)
+    mass: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 1)
     """Mass of the drone."""
 
-    gravity_vec: Array  # (N, M, 3)
+    gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Gravity vector of the drone."""
 
-    J: Array  # (N, M, 3, 3)
+    J: Array = field(metadata={CORE_NDIM_KEY: 2})  # (N, M, 3, 3)
     """Inertia matrix of the drone."""
 
-    J_inv: Array  # (N, M, 3, 3)
+    J_inv: Array = field(metadata={CORE_NDIM_KEY: 2})  # (N, M, 3, 3)
     """Inverse of the inertia matrix of the drone."""
 
-    acc_coef: Array  # (N, M, 1)
+    acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Coefficient for the acceleration."""
 
-    cmd_f_coef: Array  # (N, M, 1)
+    cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Coefficient for the collective thrust."""
 
-    rpy_coef: Array  # (N, M, 1)
+    rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Coefficient for the roll pitch yaw dynamics."""
 
-    rpy_rates_coef: Array  # (N, M, 1)
+    rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Coefficient for the roll pitch yaw rates dynamics."""
 
-    cmd_rpy_coef: Array  # (N, M, 1)
+    cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Coefficient for the roll pitch yaw command dynamics."""
 
     @staticmethod
diff --git a/crazyflow/dynamics/so_rpy_rotor/dynamics.py b/crazyflow/dynamics/so_rpy_rotor/dynamics.py
index eef3723f..ae929e5f 100644
--- a/crazyflow/dynamics/so_rpy_rotor/dynamics.py
+++ b/crazyflow/dynamics/so_rpy_rotor/dynamics.py
@@ -23,13 +23,13 @@
 import jax.numpy as jnp
 from array_api_compat import array_namespace
 from array_api_compat import device as xp_device
-from flax.struct import dataclass
+from flax.struct import dataclass, field
 from scipy.spatial.transform import Rotation as R
 
 import crazyflow.dynamics.symbols as symbols
 from crazyflow.dynamics.core import load_params, supports
 from crazyflow.dynamics.utils import rotation
-from crazyflow.utils import to_xp
+from crazyflow.utils import CORE_NDIM_KEY, to_xp
 
 if TYPE_CHECKING:
     from jax import Device
@@ -375,25 +375,25 @@ def symbolic_dynamics_euler(
 
 @dataclass
 class Params:
-    mass: Array  # (N, M, 1)
+    mass: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 1)
     """Mass of the drone."""
-    gravity_vec: Array  # (N, M, 3)
+    gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Gravity vector of the drone."""
-    J: Array  # (N, M, 3, 3)
+    J: Array = field(metadata={CORE_NDIM_KEY: 2})  # (N, M, 3, 3)
     """Inertia matrix of the drone."""
-    J_inv: Array  # (N, M, 3, 3)
+    J_inv: Array = field(metadata={CORE_NDIM_KEY: 2})  # (N, M, 3, 3)
     """Inverse of the inertia matrix of the drone."""
-    thrust_time_coef: Array  # (N, M, 1)
+    thrust_time_coef: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Rotor coefficient of the drone."""
-    acc_coef: Array  # (N, M, 1)
+    acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Acceleration coefficient of the drone."""
-    cmd_f_coef: Array  # (N, M, 1)
+    cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Collective thrust coefficient of the drone."""
-    rpy_coef: Array  # (N, M, 1)
+    rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Roll pitch yaw coefficient of the drone."""
-    rpy_rates_coef: Array  # (N, M, 1)
+    rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Roll pitch yaw rates coefficient of the drone."""
-    cmd_rpy_coef: Array  # (N, M, 1)
+    cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Roll pitch yaw command coefficient of the drone."""
 
     @staticmethod
diff --git a/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py b/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py
index c72fdc61..c39a27c0 100644
--- a/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py
+++ b/crazyflow/dynamics/so_rpy_rotor_drag/dynamics.py
@@ -25,13 +25,13 @@
 import jax.numpy as jnp
 from array_api_compat import array_namespace
 from array_api_compat import device as xp_device
-from flax.struct import dataclass
+from flax.struct import dataclass, field
 from scipy.spatial.transform import Rotation as R
 
 import crazyflow.dynamics.symbols as symbols
 from crazyflow.dynamics.core import load_params, supports
 from crazyflow.dynamics.utils import rotation
-from crazyflow.utils import to_xp
+from crazyflow.utils import CORE_NDIM_KEY, to_xp
 
 if TYPE_CHECKING:
     from jax import Device
@@ -409,27 +409,27 @@ def symbolic_dynamics_euler(
 
 @dataclass
 class Params:
-    mass: Array  # (N, M, 1)
+    mass: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 1)
     """Mass of the drone."""
-    gravity_vec: Array  # (N, M, 3)
+    gravity_vec: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Gravity vector of the drone."""
-    J: Array  # (N, M, 3, 3)
+    J: Array = field(metadata={CORE_NDIM_KEY: 2})  # (N, M, 3, 3)
     """Inertia matrix of the drone."""
-    J_inv: Array  # (N, M, 3, 3)
+    J_inv: Array = field(metadata={CORE_NDIM_KEY: 2})  # (N, M, 3, 3)
     """Inverse of the inertia matrix of the drone."""
-    thrust_time_coef: Array  # (N, M, 1)
+    thrust_time_coef: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Rotor coefficient of the drone."""
-    acc_coef: Array  # (N, M, 1)
+    acc_coef: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Acceleration coefficient of the drone."""
-    cmd_f_coef: Array  # (N, M, 1)
+    cmd_f_coef: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Collective thrust coefficient of the drone."""
-    rpy_coef: Array  # (N, M, 1)
+    rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Roll pitch yaw coefficient of the drone."""
-    rpy_rates_coef: Array  # (N, M, 1)
+    rpy_rates_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Roll pitch yaw rates coefficient of the drone."""
-    cmd_rpy_coef: Array  # (N, M, 1)
+    cmd_rpy_coef: Array = field(metadata={CORE_NDIM_KEY: 1})  # (3,)
     """Roll pitch yaw command coefficient of the drone."""
-    drag_matrix: Array  # (N, M, 3, 3)
+    drag_matrix: Array = field(metadata={CORE_NDIM_KEY: 2})  # (3, 3)
     """Linear drag coefficient matrix of the drone."""
 
     @staticmethod
diff --git a/crazyflow/sim/data.py b/crazyflow/sim/data.py
index 6872853d..42c313c1 100644
--- a/crazyflow/sim/data.py
+++ b/crazyflow/sim/data.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing
+from typing import Any
 
 import jax
 import jax.numpy as jnp
@@ -18,23 +19,24 @@
 from crazyflow.dynamics.so_rpy import Params as SoRpyParams
 from crazyflow.dynamics.so_rpy_rotor import Params as SoRpyRotorParams
 from crazyflow.dynamics.so_rpy_rotor_drag import Params as SoRpyRotorDragParams
+from crazyflow.utils import CORE_NDIM_KEY
 
 
 @dataclass
 class SimState:
-    pos: Array  # (N, M, 3)
+    pos: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Position of the drone's center of mass."""
-    quat: Array  # (N, M, 4)
+    quat: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
     """Quaternion of the drone's orientation."""
-    vel: Array  # (N, M, 3)
+    vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Velocity of the drone's center of mass in the world frame."""
-    ang_vel: Array  # (N, M, 3)
+    ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Angular velocity of the drone in the body frame."""
-    force: Array  # (N, M, 3)  # CoM force
+    force: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3) CoM force
     """Force applied to the drone's center of mass in the world frame."""
-    torque: Array  # (N, M, 3)  # CoM torque
+    torque: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3) CoM torque
     """Torque applied to the drone's center of mass in the world frame."""
-    rotor_vel: Array  # (N, M, 4)  # Motor forces along body frame z axis
+    rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4) in RPM
     """Motor forces along body frame z axis."""
 
     @staticmethod
@@ -57,15 +59,15 @@ def create(n_worlds: int, n_drones: int, device: Device) -> SimState:
 
 @dataclass
 class SimStateDeriv:
-    vel: Array  # (N, M, 3)
+    vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Derivative of the position of the drone's center of mass."""
-    ang_vel: Array  # (N, M, 3)
+    ang_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Derivative of the quaternion of the drone's orientation as angular velocity."""
-    acc: Array  # (N, M, 3)
+    acc: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Derivative of the velocity of the drone's center of mass."""
-    ang_acc: Array  # (N, M, 3)
+    ang_acc: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Derivative of the angular velocity of the drone's center of mass."""
-    rotor_acc: Array  # (N, M, 4)
+    rotor_acc: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
     """Derivative of the rotor velocity."""
 
     @staticmethod
@@ -90,7 +92,7 @@ class ControlData(typing.Protocol):
     """Control command for the drone."""
     staged_cmd: Array  # (N, M, X)
     """Staged control command for the drone."""
-    steps: Array  # (N, 1)
+    steps: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, 1)
     """Last simulation steps that the state control command was applied."""
     freq: int
     """Frequency of the state control command."""
@@ -108,7 +110,7 @@ class SimControls:
     """Attitude control data."""
     force_torque: ControlData | None
     """Force and torque control data."""
-    rotor_vel: Array  # (N, M, 4)
+    rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
     """Desired motor speed."""
 
     @staticmethod
@@ -205,19 +207,17 @@ def create(
 class SimCore:
     freq: int = field(pytree_node=False)
     """Frequency of the simulation."""
-    device: Device = field(pytree_node=False)
-    """Device of the simulation."""
-    steps: Array  # (N, 1)
+    steps: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, 1)
     """Simulation steps taken since the last reset."""
     n_worlds: int = field(pytree_node=False)
     """Number of worlds in the simulation."""
     n_drones: int = field(pytree_node=False)
     """Number of drones in the simulation."""
-    drone_mocap_ids: Array  # (M,)
+    drone_mocap_ids: Array = field(metadata={CORE_NDIM_KEY: 1})  # (M,)
     """MuJoCo mocap IDs of the drone bodies."""
-    rng_key: Array  # (N, 1)
+    rng_key: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Random number generator key for the simulation."""
-    mjx_synced: Array  # (1,)
+    mjx_synced: Array = field(metadata={CORE_NDIM_KEY: 0})  # ()
     """Whether the simulation data is synchronized with the MuJoCo mjx_data."""
 
     @staticmethod
@@ -236,7 +236,6 @@ def create(
         rng_key = jax.device_put(rng_key, device)
         return SimCore(
             freq=freq,
-            device=device,
             steps=steps,
             n_worlds=n_worlds,
             n_drones=n_drones,
@@ -258,5 +257,10 @@ class SimData:
     """Drone parameters."""
     core: SimCore
     """Core parameters of the simulation."""
-    plugins: dict[str, Array] = field(default_factory=dict)
-    """Arbitrary data for plugins to store state in the simulation."""
+    plugins: dict[str, Any] = field(default_factory=dict)
+    """Arbitrary data for plugins to store state in the simulation.
+
+    Data that is batched over worlds has to declare its dimensions without batch axes using the
+    CORE_NDIM_KEY metadata, so that resets can mask it and sharding can distribute it. Dicts have no
+    fields to declare on, so batched data has to be stored in a struct.
+    """
diff --git a/crazyflow/sim/functional.py b/crazyflow/sim/functional.py
index 354af0b7..e69a769a 100644
--- a/crazyflow/sim/functional.py
+++ b/crazyflow/sim/functional.py
@@ -6,7 +6,6 @@
 
 from crazyflow.control import Control
 from crazyflow.control.core import controllable as _controllable
-from crazyflow.utils import to_device
 
 if TYPE_CHECKING:
     from jax import Array
@@ -18,7 +17,7 @@ def state_control(data: SimData, controls: Array) -> SimData:
     """State control function."""
     assert data.controls.mode == Control.state, f"control type {data.controls.mode} not enabled"
     assert controls.shape == (data.core.n_worlds, data.core.n_drones, 13), "controls shape mismatch"
-    controls = to_device(controls, data.core.device)
+    controls = jnp.asarray(controls)
     data = data.replace(
         controls=data.controls.replace(state=data.controls.state.replace(staged_cmd=controls))
     )
@@ -36,7 +35,7 @@ def attitude_control(data: SimData, controls: Array) -> SimData:
     """
     assert data.controls.mode == Control.attitude, f"control type {data.controls.mode} not enabled"
     assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch"
-    controls = to_device(controls, data.core.device)
+    controls = jnp.asarray(controls)
     data = data.replace(
         controls=data.controls.replace(attitude=data.controls.attitude.replace(staged_cmd=controls))
     )
@@ -49,7 +48,7 @@ def force_torque_control(data: SimData, controls: Array) -> SimData:
         f"control type {data.controls.mode} not enabled"
     )
     assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch"
-    controls = to_device(controls, data.core.device)
+    controls = jnp.asarray(controls)
     data = data.replace(
         controls=data.controls.replace(
             force_torque=data.controls.force_torque.replace(staged_cmd=controls)
@@ -65,7 +64,7 @@ def rotor_vel_control(data: SimData, controls: Array) -> SimData:
     """
     assert data.controls.mode == Control.rotor_vel, f"control type {data.controls.mode} not enabled"
     assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch"
-    controls = to_device(controls, data.core.device)
+    controls = jnp.asarray(controls)
     return data.replace(controls=data.controls.replace(rotor_vel=controls))
 
 
@@ -81,7 +80,7 @@ def controllable(data: SimData) -> Array:
             control_steps = controls.force_torque.steps
             control_freq = controls.force_torque.freq
         case Control.rotor_vel:
-            return jnp.ones((data.core.n_worlds, 1), dtype=bool, device=data.core.device)
+            return jnp.ones_like(data.core.steps, dtype=bool)
         case _:
             raise NotImplementedError(f"Control mode {data.controls.mode} not implemented")
     return _controllable(data.core.steps, data.core.freq, control_steps, control_freq)
diff --git a/crazyflow/sim/sharding.py b/crazyflow/sim/sharding.py
new file mode 100644
index 00000000..2312a7c7
--- /dev/null
+++ b/crazyflow/sim/sharding.py
@@ -0,0 +1,73 @@
+"""Distribution of crazyflow over multiple devices.
+
+Sharding arrays along the world axis allows efficient parallelization of crazyflow across multiple
+devices. The data itself specifies which arrays carry the world axis. See
+[world_mask][crazyflow.utils.world_mask].
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import jax
+from jax.sharding import AxisType, NamedSharding, PartitionSpec
+
+from crazyflow.utils import world_mask
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from jax import Device
+    from jax.sharding import Mesh
+
+    from crazyflow.sim.data import SimData
+
+WORLD_AXIS = "worlds"
+"""Name of the mesh axis that the worlds are distributed over."""
+
+
+def world_mesh(devices: Sequence[Device]) -> Mesh:
+    """Create a mesh that distributes the worlds over the devices.
+
+    The mesh uses automatic axis types. Explicit axis types cause problems in SciPy.
+
+    Args:
+        devices: Devices to distribute the worlds over.
+
+    Returns:
+        A one-dimensional mesh over the world axis.
+    """
+    axes, names = (len(devices),), (WORLD_AXIS,)
+    return jax.make_mesh(axes, names, axis_types=(AxisType.Auto,), devices=devices)
+
+
+def placement(data: SimData, mesh: Mesh) -> SimData:
+    """Build the placement that distributes the worlds of the simulation data over a mesh.
+
+    Note:
+        Pass the placement to `jax.device_put`, or modify it first to place the data by hand.
+
+    Args:
+        data: Simulation data to place.
+        mesh: Mesh to distribute the worlds over, as built by
+            [world_mesh][crazyflow.sim.sharding.world_mesh].
+
+    Returns:
+        A pytree of shardings matching the simulation data.
+    """
+    world = NamedSharding(mesh, PartitionSpec(WORLD_AXIS))
+    replicated = NamedSharding(mesh, PartitionSpec())
+    return jax.tree.map(lambda indexed: world if indexed else replicated, world_mask(data))
+
+
+def shard(data: SimData, mesh: Mesh) -> SimData:
+    """Distribute the worlds of the simulation data over a mesh.
+
+    Args:
+        data: Simulation data to place.
+        mesh: Mesh to distribute the worlds over.
+
+    Returns:
+        The placed simulation data.
+    """
+    return jax.device_put(data, placement(data, mesh))
diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py
index a4abc884..fe955a8b 100644
--- a/crazyflow/sim/sim.py
+++ b/crazyflow/sim/sim.py
@@ -31,9 +31,11 @@
 from crazyflow.sim.data import SimControls, SimCore, SimData, SimParams, SimState, SimStateDeriv
 from crazyflow.sim.integration import Integrator, euler, rk4, symplectic_euler
 from crazyflow.sim.pipeline import append_fn
-from crazyflow.utils import grid_2d, pytree_replace
+from crazyflow.sim.sharding import placement
+from crazyflow.utils import grid_2d, pytree_replace, world_mask
 
 if TYPE_CHECKING:
+    from jax.sharding import Mesh
     from mujoco.mjx import Data, Model
     from numpy.typing import NDArray
 
@@ -430,6 +432,20 @@ def build_data(self) -> SimData:
         )
         return self.data
 
+    def shard(self, mesh: Mesh) -> SimData:
+        """Distribute the data and default data over a mesh along the world axis.
+
+        Args:
+            mesh: Mesh to distribute the worlds over, as built by
+                [world_mesh][crazyflow.sim.sharding.world_mesh].
+
+        Returns:
+            The placed simulation data.
+        """
+        self.data = jax.device_put(self.data, placement(self.data, mesh))
+        self.default_data = jax.device_put(self.default_data, placement(self.default_data, mesh))
+        return self.data
+
     def build_default_data(self) -> SimData:
         """Initialize the default data for the simulation.
 
@@ -601,8 +617,14 @@ def select_integrate_fn(
 
 
 def reset(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
-    """Reset the simulation data to the default data for the worlds specified by the mask."""
-    return pytree_replace(data, default_data, mask)  # Does not overwrite rng_key
+    """Reset the simulation data to the default data for the worlds specified by the mask.
+
+    Without a mask, the full data is restored. The mask selects along the world axis, so it only
+    restores per-world arrays. The rng key is never restored.
+    """
+    if mask is None:
+        return default_data.replace(core=default_data.core.replace(rng_key=data.core.rng_key))
+    return pytree_replace(data, default_data, world_mask(data), mask)
 
 
 def increment_steps(data: SimData) -> SimData:
diff --git a/crazyflow/utils.py b/crazyflow/utils.py
index 5136444d..e0f97f7a 100644
--- a/crazyflow/utils.py
+++ b/crazyflow/utils.py
@@ -3,6 +3,7 @@
 import inspect
 import os
 from collections.abc import Mapping
+from dataclasses import fields, is_dataclass
 from functools import partial
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, ParamSpec, TypeVar
@@ -15,6 +16,9 @@
 if TYPE_CHECKING:
     from types import ModuleType
 
+CORE_NDIM_KEY = "core_ndim"
+"""Metadata key for the number of dimensions a dataclass field has without any batch axes."""
+
 
 def grid_2d(n: int, spacing: float = 1.0, center: Array | None = None) -> Array:
     """Generate a 2D grid of points."""
@@ -32,21 +36,64 @@ def grid_2d(n: int, spacing: float = 1.0, center: Array | None = None) -> Array:
 R = TypeVar("R")
 
 
-def pytree_replace(tree: T, new_tree: T, mask: Array | None = None) -> T:
-    """Overwrite elements of a PyTree with values from another PyTree filtered by a mask.
+def world_mask(tree: T) -> T:
+    """Flag the arrays of a PyTree that are indexed by world.
+
+    Fields declare their number of dimensions without batch axes with the CORE_NDIM_KEY metadata.
+    An array carries a world axis when its ndim exceeds that. Arrays without the metadata are
+    treated as not batched.
+
+    Args:
+        tree: PyTree of arrays and structs to classify.
+
+    Returns:
+        A PyTree of booleans matching the input.
+    """
+    return _flag(tree)
+
+
+def _flag(node: Any) -> Any:
+    """Flag the leaves of a node, recursing through the declared structure of the data."""
+    if not _declares_fields(node):
+        # Any other node. Nested structs are handed back to us via is_leaf so that we can recurse
+        return jax.tree.map(
+            lambda x: _flag(x) if _declares_fields(x) else False, node, is_leaf=_declares_fields
+        )
+    flags = {}
+    for f in (f for f in fields(node) if f.metadata.get("pytree_node", True)):
+        value, core_ndim = getattr(node, f.name), f.metadata.get(CORE_NDIM_KEY)
+        if core_ndim is None:  # Structs and containers classify their own contents
+            flags[f.name] = _flag(value)
+        elif _declares_fields(value):
+            name = f"{type(node).__name__}.{f.name}"
+            raise ValueError(f"{name} holds a struct, which declares its own fields")
+        else:
+            flags[f.name] = jax.tree.map(lambda x, core=core_ndim: x.ndim > core, value)
+    return node.replace(**flags)
+
+
+def _declares_fields(node: Any) -> bool:
+    """Check whether a node is a dataclass that jax traverses."""
+    return is_dataclass(node) and not jax.tree_util.all_leaves([node])
+
 
-    The mask indicates which elements of the leaf arrays to overwrite with new values, and which
-    ones to leave unchanged.
+def pytree_replace(tree: T, new_tree: T, batched: T, mask: Array | None = None) -> T:
+    """Overwrite batched leaves of a PyTree with values from another PyTree filtered by a mask.
+
+    Args:
+        tree: PyTree to overwrite.
+        new_tree: PyTree to take the new values from.
+        batched: PyTree of booleans flagging the leaves that carry the batch axis.
+        mask: Boolean array matching the leading axis of the batched leaves.
     """
 
-    def _replace(x: Array, y: Array) -> Array:
-        """Replace broadcastable leaves in tree.map."""
-        # Do not replace leaves that are not broadcastable
-        if x.ndim < 1 or (mask is not None and mask.shape[0] != x.shape[0]):
+    def _replace(x: Array, y: Array, batched: bool) -> Array:
+        """Replace batched leaves in tree.map."""
+        if not batched:
             return x
         return jnp.where(broadcast_mask(mask, x.shape), y, x)
 
-    return jax.tree.map(_replace, tree, new_tree)
+    return jax.tree.map(_replace, tree, new_tree, batched)
 
 
 def leaf_replace(tree: T, mask: Array | None = None, **kwargs: dict[str, Array]) -> T:
@@ -71,12 +118,6 @@ def broadcast_mask(mask: Array | None, shape: tuple[int, ...]) -> Array:
     return mask.reshape(*mask.shape, *[1] * (len(shape) - mask.ndim))
 
 
-@partial(jax.jit, static_argnames="device")
-def to_device(data: Array, device: str) -> Array:
-    """Turn an array into a jax array on the specified device."""
-    return jnp.array(data, device=device)
-
-
 def enable_cache(
     cache_path: Path | None = None,
     min_entry_size_bytes: int = -1,
diff --git a/docs/api/index.md b/docs/api/index.md
index 6463c252..8c602fab 100644
--- a/docs/api/index.md
+++ b/docs/api/index.md
@@ -10,6 +10,7 @@ This section is auto-generated from the Crazyflow source code using [mkdocstring
 | `crazyflow.sim.pipeline` | `OrderedDict`-based pipeline helpers (`append_fn`, `insert_fn_before`, `replace_fn`, etc.) |
 | `crazyflow.sim.data` | `SimData`, `SimState`, `SimControls`, `SimParams`, `SimCore` pytrees |
 | `crazyflow.sim.functional` | Pure functional control API for use inside `jax.jit` |
+| `crazyflow.sim.sharding` | Placement of the simulation data on multiple devices |
 | `crazyflow.dynamics` | `Dynamics` enum and dynamics implementations |
 | `crazyflow.sim.integration` | `Integrator` enum, Euler, RK4, and symplectic Euler |
 | `crazyflow.sim.sensors` | Raycast depth (`sensors.depth`) and gaussian splat RGB (`sensors.splat`) sensors |
diff --git a/docs/examples/index.md b/docs/examples/index.md
index 76c876cb..ab5322e2 100644
--- a/docs/examples/index.md
+++ b/docs/examples/index.md
@@ -51,6 +51,17 @@ Because the simulator is built entirely from JAX operations, `jax.grad` can diff
 
 ---
 
+## Sharding across devices
+
+We can also distribute the simulation across devices. The host backend is asked for four logical devices. We then compare this against running on a single device. See [Sharding](../user-guide/sharding.md) for the details.
+
+<!-- notest: imported script, covered by tests/integration/test_examples.py -->
+```{ .python notest }
+--8<-- "examples/jax/sharding.py"
+```
+
+---
+
 ## Domain randomization
 
 Randomizing mass and inertia through the reset pipeline. An optional mask limits randomization to selected worlds.
diff --git a/docs/gen_ref_pages.py b/docs/gen_ref_pages.py
index 52d418ac..88e40cbf 100644
--- a/docs/gen_ref_pages.py
+++ b/docs/gen_ref_pages.py
@@ -53,6 +53,7 @@
     * [sim.sensors](crazyflow/sim/sensors/index.md)
     * [sim.sensors.depth](crazyflow/sim/sensors/depth.md)
     * [sim.sensors.splat](crazyflow/sim/sensors/splat.md)
+    * [sim.sharding](crazyflow/sim/sharding.md)
     * [sim.sim](crazyflow/sim/sim.md)
     * [sim.splat](crazyflow/sim/splat.md)
     * [sim.visualize](crazyflow/sim/visualize.md)
diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md
index 21663c91..99fbe308 100644
--- a/docs/user-guide/index.md
+++ b/docs/user-guide/index.md
@@ -8,6 +8,7 @@ In-depth documentation for every part of the simulator.
 - [Dynamics](dynamics/index.md) — first-principles vs. fitted dynamics, when to use each
 - [Control Modes](control/index.md) — state, attitude, force/torque, and rotor velocity control
 - [Pipelines](pipelines.md) — composable step and reset pipelines, randomization, and disturbances
+- [The world axis](world-axis.md) — which arrays are batched over worlds, and what resets and sharding do with them
 - [Visualization](visualization.md) — rendering modes, cameras, raycasting, and materials
 - [MuJoCo Integration](mujoco.md) — MJCF scene construction, adding objects, and sync internals
 - [Gymnasium Environments](gymnasium-envs.md) — vectorized environments for RL training
diff --git a/docs/user-guide/oo-api.md b/docs/user-guide/oo-api.md
index b8389bc3..c13fb36e 100644
--- a/docs/user-guide/oo-api.md
+++ b/docs/user-guide/oo-api.md
@@ -138,6 +138,8 @@ Passing more steps to a single `step(n_steps)` call is more efficient than multi
 
 `sim.reset()` reinitialises all worlds to their default state. Pass a boolean mask of shape `(n_worlds,)` to reset only selected worlds: `True` resets that world, `False` leaves it unchanged. This is useful in RL training loops where episodes end at different times.
 
+A full reset restores everything except the rng key. A mask selects along the world axis, so it only restores per-world arrays, and leaves parameters shared by all worlds untouched. See [The world axis](world-axis.md) for more details.
+
 ```python
 import numpy as np
 from crazyflow.sim import Sim
diff --git a/docs/user-guide/sharding.md b/docs/user-guide/sharding.md
new file mode 100644
index 00000000..f5873612
--- /dev/null
+++ b/docs/user-guide/sharding.md
@@ -0,0 +1,55 @@
+# Sharding
+
+Crazyflow can be sharded across devices for even more performance. We partition the simulation data along the world axis and replicate shared data. `sim.shard` places `sim.data` and `sim.default_data` on multiple devices, and subsequent runs will use all devices to compute the simulation step.
+
+!!! note
+    Workloads that do not make use of sharding are unaffected by replication etc.
+
+```python
+import jax
+
+from crazyflow.sim import Sim
+from crazyflow.sim.sharding import world_mesh
+
+devices = jax.devices()
+sim = Sim(n_worlds=4 * len(devices))
+sim.shard(world_mesh(devices))
+sim.step()
+
+assert sim.data.states.pos.sharding.spec == jax.sharding.PartitionSpec("worlds")
+assert sim.data.params.gravity_vec.sharding.spec == jax.sharding.PartitionSpec()
+```
+
+The number of worlds has to be divisible by the number of devices. Whether an array is batched over worlds is determined from its shape and the core ndim its field declares, see [The world axis](world-axis.md). Data shared across worlds, such as the gravity vector above, is replicated on every device.
+
+!!! warning
+    Sharding relies on JAX's automatic sharding mode, which `world_mesh` requests for you. `jax.make_mesh` defaults to explicit sharding, under which the Mellinger controller fails inside SciPy's rotation backend.
+
+## Functional API
+
+`shard` places a `SimData` without going through the `Sim` object:
+
+```{ .python continuation }
+from crazyflow.sim.sharding import shard
+
+mesh = world_mesh(devices)
+data, default_data = shard(sim.data, mesh), shard(sim.default_data, mesh)
+```
+
+`placement` returns the shardings that `shard` applies, as a pytree mirroring the simulation data. Modify it to place the data by hand:
+
+```{ .python continuation }
+from crazyflow.sim.sharding import placement
+
+data = jax.device_put(sim.data, placement(sim.data, mesh))
+```
+
+## Plugin data
+
+Bare arrays in `sim.data.plugins` are replicated, since a dict has no fields to declare a core ndim on. Store per-world plugin state in a struct that declares one, as shown in [The world axis](world-axis.md#your-own-state).
+
+## Next steps
+
+- [The world axis](world-axis.md) — declaring which arrays are batched over worlds
+- [Functional API](functional-api.md) — composing the pure step, reset and control functions
+- [Batching & domain randomization](dynamics/batching.md) — varying the drone parameters per world
diff --git a/docs/user-guide/world-axis.md b/docs/user-guide/world-axis.md
new file mode 100644
index 00000000..e97f97cf
--- /dev/null
+++ b/docs/user-guide/world-axis.md
@@ -0,0 +1,87 @@
+# The world axis
+
+The simulation is batched over worlds, so most arrays in `sim.data` carry a leading `n_worlds` axis. Resetting selected worlds and distributing worlds across devices both need to know which arrays those are. Here we explain how that is declared and what depends on it.
+
+## How it is declared
+
+A world axis is an ordinary axis, so no shape on its own tells you whether it indexes worlds. With three worlds, `params.drag_matrix` of shape `(3, 3)` and `params.gravity_vec` of shape `(3,)` both look world-batched while neither is.
+
+Each field therefore declares its number of dimensions without batch axes with the `CORE_NDIM_KEY` metadata. `states.pos` is a vector, so its core ndim is 1, and `params.drag_matrix` is a matrix, so its core ndim is 2. An array carries a world axis exactly when its own `ndim` exceeds the core ndim of its field:
+
+```python
+from crazyflow.sim import Sim
+from crazyflow.utils import world_mask
+
+sim = Sim(n_worlds=3)
+mask = world_mask(sim.data)
+
+assert mask.states.pos  # (n_worlds, n_drones, 3), indexed by world
+assert mask.params.mass  # (n_worlds, n_drones, 1), one mass per drone
+assert not mask.params.gravity_vec  # (3,), shared by all worlds
+assert not mask.params.drag_matrix  # (3, 3), shared by all worlds
+```
+
+`world_mask` returns a pytree of booleans matching the data, with one flag per array. Arrays whose field declares no core ndim are shared by all worlds.
+
+## Resets
+
+`sim.reset()` restores the whole simulation, both the world-indexed and the shared arrays, with the rng key as the only exception.
+
+`sim.reset(mask)` is the constrained form. It selects worlds along the world axis, so it can only restore arrays that have one. A shared array cannot be restored for some worlds and not others, so a masked reset leaves it untouched:
+
+```{ .python continuation }
+import jax.numpy as jnp
+
+gravity = jnp.array([1.0, 2.0, 3.0])
+sim.data = sim.data.replace(params=sim.data.params.replace(gravity_vec=gravity))
+sim.reset(jnp.array([True, False, False]))
+
+assert jnp.array_equal(sim.data.params.gravity_vec, gravity)  # Shared, so left alone
+assert jnp.all(sim.data.states.pos[0] == 0.0)  # World 0 was reset
+```
+
+To keep a change to a shared parameter across masked resets, write it into `sim.default_data` as well by calling `sim.build_default_data()`. See [Stepping and resetting](oo-api.md#stepping-and-resetting) for the reset API.
+
+## Sharding
+
+Distributing the simulation partitions the world axis of the arrays that have one and replicates the rest, using the same declarations. See [Sharding](sharding.md) for the mesh and placement API.
+
+## Randomizing a shared parameter
+
+A parameter that is shared by default gains a world axis the moment you randomize it per world, and its `ndim` then exceeds the declared core ndim:
+
+```{ .python continuation }
+n = sim.n_worlds
+scale = jnp.arange(1, n + 1)[:, None, None, None]
+drag = jnp.broadcast_to(sim.data.params.drag_matrix, (n, 1, 3, 3)) * scale
+sim.data = sim.data.replace(params=sim.data.params.replace(drag_matrix=drag))
+sim.build_default_data()  # Keep the randomized shape across full resets
+
+assert world_mask(sim.data).params.drag_matrix
+```
+
+From there it behaves like any other world-indexed array. A masked reset restores it for the masked worlds, and sharding partitions it instead of replicating it. Setting it back to a `(3, 3)` array makes it shared again.
+
+## Your own state
+
+Anything you add to `sim.data.plugins` follows the same rule. A `dict` has no fields, so bare arrays are shared, which means a masked reset leaves them alone and sharding replicates them. If you need per-world state, declare it on a struct:
+
+```python
+import flax.struct
+from jax import Array
+
+from crazyflow.utils import CORE_NDIM_KEY
+
+
+@flax.struct.dataclass
+class PluginState:
+    buffer: Array = flax.struct.field(metadata={CORE_NDIM_KEY: 1})
+    lookup: Array = None
+```
+
+If you store this under `sim.data.plugins`, a `buffer` of shape `(n_worlds, 3)` exceeds its core ndim of 1, so it is restored for the masked worlds and partitioned when the simulation is sharded. A `lookup` of any shape declares no core ndim, so it is only restored by a full reset and replicated across devices.
+
+## Next steps
+
+- [Sharding](sharding.md) — distributing the worlds over several devices
+- [Pipelines](pipelines.md) — customising what runs on a step and on a reset
diff --git a/examples/jax/sharding.py b/examples/jax/sharding.py
new file mode 100644
index 00000000..9630fe40
--- /dev/null
+++ b/examples/jax/sharding.py
@@ -0,0 +1,39 @@
+"""Example on how to shard the simulation across devices."""
+
+import os
+
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"
+
+import jax
+import numpy as np
+
+from crazyflow.sim import Sim
+from crazyflow.sim.sharding import world_mesh
+
+
+def main():
+    devices = jax.devices("cpu")
+    sim = Sim(n_worlds=4 * len(devices), device="cpu")  # Worlds must be divisible by n_devices
+    sim.step(10)
+    single_device_pos = np.asarray(sim.data.states.pos)
+
+    sim.reset()
+    sim.shard(world_mesh(devices))
+    sim.step(10)
+
+    pos, gravity = sim.data.states.pos, sim.data.params.gravity_vec
+    print(f"Distributing {sim.n_worlds} worlds over {len(devices)} devices")
+    # Arrays that carry a world axis are partitioned. Shared params are replicated on every device
+    print("  Placement")
+    print(f"    {'position':<{10}}  {str(pos.shape):<{10}}  {pos.sharding.spec}")
+    print(f"    {'gravity':<{10}}  {str(gravity.shape):<{10}}  {gravity.sharding.spec}")
+    print("  Shards of states.pos")
+    for shard in pos.addressable_shards:
+        print(f"    {shard.device}  {shard.data.shape[0]} worlds")
+
+    assert np.allclose(np.asarray(pos), single_device_pos, atol=1e-6)
+    sim.close()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/properdocs.yml b/properdocs.yml
index 04f557a6..3463a604 100644
--- a/properdocs.yml
+++ b/properdocs.yml
@@ -70,6 +70,8 @@ nav:
       - Batching: user-guide/control/batching.md
       - JIT compilation: user-guide/control/jit.md
     - Pipelines: user-guide/pipelines.md
+    - The world axis: user-guide/world-axis.md
+    - Sharding: user-guide/sharding.md
     - Visualization: user-guide/visualization.md
     - Gaussian Splat Rendering: user-guide/splats.md
     - MuJoCo Integration: user-guide/mujoco.md
diff --git a/tests/conftest.py b/tests/conftest.py
index e0611592..7c88d620 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,6 +2,8 @@
 
 os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
 os.environ["SCIPY_ARRAY_API"] = "1"
+# We need multiple devices for sharding tests
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=2"
 
 import jax
 import pytest
diff --git a/tests/unit/test_plugins.py b/tests/unit/test_plugins.py
index f31950aa..1a27d76b 100644
--- a/tests/unit/test_plugins.py
+++ b/tests/unit/test_plugins.py
@@ -1,34 +1,52 @@
 """Unit tests for the simulation plugin system.
 
 Plugins are callables of the form ``fn(data: SimData) -> SimData`` that are inserted into
-``sim.step_pipeline``. They can store arbitrary state in ``sim.data.plugins`` (a dict of JAX
-arrays). Tests use a simple per-world step counter as a canonical plugin.
+``sim.step_pipeline``. They can store arbitrary state in ``sim.data.plugins``. State that is batched
+over worlds has to declare its dimensions without batch axes so that resets can mask it.
 """
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+import flax
 import jax.numpy as jnp
 import pytest
+from flax.struct import field
 
 from crazyflow.sim import Sim
 from crazyflow.sim.pipeline import append_fn
+from crazyflow.utils import CORE_NDIM_KEY
 
 if TYPE_CHECKING:
+    from jax import Array
+
     from crazyflow.sim.data import SimData
 
 
+@flax.struct.dataclass
+class Counter:
+    """Per-world plugin state."""
+
+    count: Array = field(metadata={CORE_NDIM_KEY: 1})
+
+    @staticmethod
+    def create(n_worlds: int) -> Counter:
+        """Create a zeroed counter for all worlds."""
+        return Counter(jnp.zeros((n_worlds, 1), jnp.int32))
+
+
 def counter_plugin(data: SimData) -> SimData:
     """Increment plugins["counter"] by 1 on every simulation step."""
-    return data.replace(plugins=data.plugins | {"counter": data.plugins["counter"] + 1})
+    counter = data.plugins["counter"]
+    return data.replace(plugins=data.plugins | {"counter": Counter(counter.count + 1)})
 
 
 def accumulater_plugin(data: SimData) -> SimData:
     """Accumulate plugins["counter"] on every simulation step."""
+    accumulated, counter = data.plugins["accumulated"], data.plugins["counter"]
     return data.replace(
-        plugins=data.plugins
-        | {"accumulated": data.plugins["accumulated"] + data.plugins["counter"]}
+        plugins=data.plugins | {"accumulated": Counter(accumulated.count + counter.count)}
     )
 
 
@@ -56,12 +74,12 @@ def test_builds():
 def test_plugin_data_changes(n_worlds: int):
     """Plugin data changes across simulation steps."""
     sim = Sim(n_worlds=n_worlds)
-    sim.data = sim.data.replace(plugins={"counter": jnp.zeros((n_worlds, 1), dtype=jnp.int32)})
+    sim.data = sim.data.replace(plugins={"counter": Counter.create(n_worlds)})
     append_fn(sim.step_pipeline, counter_plugin)
     sim.build_step_fn()
     n_steps = 7
     sim.step(n_steps)
-    assert jnp.all(sim.data.plugins["counter"] == n_steps), (
+    assert jnp.all(sim.data.plugins["counter"].count == n_steps), (
         f"Expected counter={n_steps}, got {sim.data.plugins['counter']}"
     )
     sim.close()
@@ -71,14 +89,14 @@ def test_plugin_data_changes(n_worlds: int):
 def test_plugin_resets():
     """sim.reset() restores the counter to its default value (0)."""
     sim = Sim()
-    sim.data = sim.data.replace(plugins={"counter": jnp.zeros((1, 1), dtype=jnp.int32)})
+    sim.data = sim.data.replace(plugins={"counter": Counter.create(1)})
     append_fn(sim.step_pipeline, counter_plugin)
     sim.build_default_data()
     sim.build_step_fn()
     sim.step(10)
-    assert jnp.all(sim.data.plugins["counter"] == 10), "Precondition: counter should be 10"
+    assert jnp.all(sim.data.plugins["counter"].count == 10), "Counter should be 10"
     sim.reset()
-    assert jnp.all(sim.data.plugins["counter"] == 0), "Counter should be 0 after full reset"
+    assert jnp.all(sim.data.plugins["counter"].count == 0), "Counter should be 0 after full reset"
     sim.close()
 
 
@@ -87,17 +105,17 @@ def test_plugin_masked_reset():
     """Masked reset only resets the plugin data for selected worlds."""
     n_worlds = 2
     sim = Sim(n_worlds=n_worlds)
-    sim.data = sim.data.replace(plugins={"counter": jnp.zeros((n_worlds, 1), dtype=jnp.int32)})
+    sim.data = sim.data.replace(plugins={"counter": Counter.create(n_worlds)})
     append_fn(sim.step_pipeline, counter_plugin)
     sim.build_default_data()
     sim.build_step_fn()
     sim.step(10)
-    assert jnp.all(sim.data.plugins["counter"] == 10), "Precondition: both counters should be 10"
+    assert jnp.all(sim.data.plugins["counter"].count == 10), "Both counters should be 10"
 
     mask = jnp.array([True, False])  # reset world 0 only
     sim.reset(mask)
-    assert jnp.all(sim.data.plugins["counter"][0] == 0), "World 0 counter must be reset to 0"
-    assert jnp.all(sim.data.plugins["counter"][1] == 10), "World 1 counter must remain at 10"
+    assert jnp.all(sim.data.plugins["counter"].count[0] == 0), "World 0 counter must be reset to 0"
+    assert jnp.all(sim.data.plugins["counter"].count[1] == 10), "World 1 counter must remain at 10"
     sim.close()
 
 
@@ -106,10 +124,7 @@ def test_chained_plugins():
     """Two chained plugins produce different results depending on their order in the pipeline."""
     sim = Sim()
     sim.data = sim.data.replace(
-        plugins={
-            "counter": jnp.zeros((1, 1), dtype=jnp.int32),
-            "accumulated": jnp.zeros((1, 1), dtype=jnp.int32),
-        }
+        plugins={"counter": Counter.create(1), "accumulated": Counter.create(1)}
     )
     append_fn(sim.step_pipeline, counter_plugin)
     append_fn(sim.step_pipeline, accumulater_plugin)
@@ -117,10 +132,10 @@ def test_chained_plugins():
     sim.build_step_fn()
     n_steps = 3
     sim.step(n_steps)
-    assert int(sim.data.plugins["counter"][0, 0]) == n_steps, (
+    assert int(sim.data.plugins["counter"].count[0, 0]) == n_steps, (
         f"Expected counter={n_steps}, got {sim.data.plugins['counter']}"
     )
-    assert int(sim.data.plugins["accumulated"][0, 0]) == sum(range(1, n_steps + 1)), (
+    assert int(sim.data.plugins["accumulated"].count[0, 0]) == sum(range(1, n_steps + 1)), (
         f"Expected accumulated={sum(range(1, n_steps + 1))}, got {sim.data.plugins['accumulated']}"
     )
     sim.close()
diff --git a/tests/unit/test_sharding.py b/tests/unit/test_sharding.py
new file mode 100644
index 00000000..55488c81
--- /dev/null
+++ b/tests/unit/test_sharding.py
@@ -0,0 +1,89 @@
+"""Unit tests for distributing the simulation across devices."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from conftest import available_backends
+from jax.sharding import PartitionSpec
+
+from crazyflow.control import Control
+from crazyflow.sim import Sim
+from crazyflow.sim.functional import attitude_control
+from crazyflow.sim.sharding import placement, shard, world_mesh
+
+# Sharding tests require at least 2 devices
+multi_device = [
+    pytest.param(
+        platform,
+        marks=pytest.mark.skipif(
+            len(jax.devices(platform)) < 2, reason=f"requires at least 2 {platform} devices"
+        ),
+    )
+    for platform in available_backends()
+]
+
+
+def assert_sharding(x: Any, sharding: Any):
+    assert x.sharding.is_equivalent_to(sharding, x.ndim), f"{x.sharding} != {sharding}"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("platform", multi_device)
+def test_placement(platform: str):
+    devices = jax.devices(platform)
+    sim = Sim(n_worlds=2 * len(devices), device=platform)
+    sim.shard(world_mesh(devices))
+    assert sim.data.states.pos.sharding.spec == PartitionSpec("worlds")
+    assert sim.data.params.mass.sharding.spec == PartitionSpec("worlds")
+    assert sim.data.params.gravity_vec.sharding.spec == PartitionSpec()
+    assert sim.data.params.rotor_dyn_coef.sharding.spec == PartitionSpec()
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("platform", multi_device)
+def test_sharded_oop_api(platform: str):
+    devices = jax.devices(platform)
+    sim = Sim(n_worlds=2 * len(devices), device=platform, control=Control.attitude)
+    sim.shard(world_mesh(devices))
+    expected = placement(sim.data, world_mesh(jax.devices(platform)))
+    jax.tree.map(assert_sharding, sim.default_data, expected)
+    sim.attitude_control(jnp.zeros((sim.n_worlds, 1, 4)))
+    sim.step()
+    jax.tree.map(assert_sharding, sim.data, expected)
+    sim.reset()
+    jax.tree.map(assert_sharding, sim.data, expected)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("platform", multi_device)
+def test_sharded_functional_api(platform: str):
+    devices = jax.devices(platform)
+    sim = Sim(n_worlds=2 * len(devices), control=Control.attitude, device=platform)
+    step_fn, reset_fn = sim.build_step_fn(), sim.build_reset_fn()
+    mesh = world_mesh(devices)
+    data, default_data = shard(sim.data, mesh), shard(sim.default_data, mesh)
+    expected = placement(sim.data, mesh)
+    data = attitude_control(data, jnp.zeros((sim.n_worlds, 1, 4)))
+    data = step_fn(data)
+    jax.tree.map(assert_sharding, data, expected)
+    data = reset_fn(data, default_data)
+    jax.tree.map(assert_sharding, data, expected)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("platform", multi_device)
+def test_sharded_step_values(platform: str):
+    # Sharding must not change the simulation results
+    devices = jax.devices(platform)
+    sim = Sim(n_worlds=2 * len(devices), device=platform)
+    sim.step(10)
+    pos = np.asarray(sim.data.states.pos)
+    sim.reset()
+    sim.shard(world_mesh(devices))
+    sim.step(10)
+    assert np.allclose(np.asarray(sim.data.states.pos), pos, atol=1e-6)
diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py
index 7d91a382..3a6dde58 100644
--- a/tests/unit/test_sim.py
+++ b/tests/unit/test_sim.py
@@ -530,7 +530,7 @@ def assert_committed(obj0: Array | Any, path: str = "data"):
         elif isinstance(obj0, (list, tuple)):  # Handle sequences
             for i, item0 in enumerate(obj0):
                 assert_committed(item0, f"{path}[{i}]")
-        elif isinstance(obj0, type(sim.data.core.device)):  # Device objects
+        elif isinstance(obj0, type(sim.device)):  # Device objects
             pass  # Devices themselves don't have committed attribute
         elif isinstance(obj0, dict):
             for key, value0 in obj0.items():
@@ -640,3 +640,26 @@ def test_fused_model(device: str, drone: str):
     sim.reset()
     sim.step(1)
     sim.close()
+
+
+@pytest.mark.unit
+def test_partial_reset_keeps_shared_arrays():
+    # A partial reset with a world mask must leave shared arrays intact
+    sim = Sim(n_worlds=3)
+    gravity = jnp.array([1.0, 2.0, 3.0])
+    sim.data = sim.data.replace(params=sim.data.params.replace(gravity_vec=gravity))
+    sim.reset(jnp.array([True, False, False]))
+    assert jnp.array_equal(sim.data.params.gravity_vec, gravity)
+
+
+@pytest.mark.unit
+def test_full_reset_restores_shared_arrays():
+    sim = Sim(n_worlds=3)
+    default_gravity, rng_key = sim.default_data.params.gravity_vec, sim.data.core.rng_key
+    sim.data = sim.data.replace(
+        params=sim.data.params.replace(gravity_vec=jnp.array([1.0, 2.0, 3.0]))
+    )
+    sim.reset()
+    assert jnp.array_equal(sim.data.params.gravity_vec, default_gravity)
+    # The random key is the only thing that does not reset
+    assert jnp.array_equal(jax.random.key_data(sim.data.core.rng_key), jax.random.key_data(rng_key))
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index bb83d018..3ae49075 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -1,7 +1,20 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import flax
 import jax
+import jax.numpy as jnp
 import pytest
+from flax.struct import field
 
-from crazyflow.utils import enable_cache
+from crazyflow.control import Control
+from crazyflow.sim import Sim
+from crazyflow.utils import CORE_NDIM_KEY, enable_cache, world_mask
+
+if TYPE_CHECKING:
+    from jax import Array
 
 
 @pytest.mark.unit
@@ -47,3 +60,87 @@ def test_enable_cache(enable_xla: bool):
             jax.config.update("jax_persistent_cache_min_compile_time_secs", orig_min_time)
         if orig_xla is not None:
             jax.config.update("jax_persistent_cache_enable_xla_caches", orig_xla)
+
+
+@pytest.mark.unit
+def test_world_mask():
+    # Check that the mask correctly handles cases that previously failed under shape-based detection
+    mask = world_mask(Sim(n_worlds=3, control=Control.attitude).data)
+    assert mask.states.pos
+    assert mask.states_deriv.acc
+    assert mask.core.steps
+    assert mask.params.mass
+    assert mask.controls.attitude.cmd
+    assert not mask.params.gravity_vec
+    assert not mask.params.drag_matrix
+    assert not mask.core.drone_mocap_ids
+    assert not mask.controls.attitude.params["kR"]
+
+
+@pytest.mark.unit
+def test_world_mask_covers_batched_arrays():
+    # Make a shape-based check that we catch all world-batched arrays using a unique world length
+    sim = Sim(n_worlds=7, control=Control.attitude)
+    paths, leaves = jax.tree.flatten_with_path(sim.data)[0], jax.tree.leaves(world_mask(sim.data))
+    for (path, x), batched in zip(paths, leaves):
+        world_axis = x.ndim >= 2 and x.shape[0] == 7
+        assert batched == world_axis, f"{jax.tree_util.keystr(path)}: {x.shape=} is {batched=}"
+
+
+@pytest.mark.unit
+def test_world_mask_user_dataclass():
+    @flax.struct.dataclass
+    class PluginState:
+        buffer: Array = field(metadata={CORE_NDIM_KEY: 1})
+        constant: Array = None
+
+    state = PluginState(buffer=jnp.zeros((3, 2)), constant=jnp.zeros(3))
+    mask = world_mask({"state": state, "array": jnp.zeros((3, 2))})
+    assert mask["state"].buffer
+    assert not mask["state"].constant
+    assert not mask["array"]  # Arrays that no struct declares are shared
+
+
+@pytest.mark.unit
+def test_world_mask_matches_structure():
+    # The mask has to match the tree leaf for leaf, whatever container holds the arrays
+    @flax.struct.dataclass
+    class PluginState:  # With world data
+        buffer: Array = field(metadata={CORE_NDIM_KEY: 1})
+
+    @dataclass
+    class Opaque:  # State that is treated as leaf
+        a: int
+
+    array = jnp.zeros((3, 2))
+    tree = {
+        "containers": (array, [array], {"x": array}),
+        "nested": (PluginState(buffer=array),),
+        "opaque": Opaque(1),
+    }
+    mask = world_mask(tree)
+    assert jax.tree.structure(mask) == jax.tree.structure(tree)
+    assert mask["nested"][0].buffer  # A declaration survives inside a container
+
+
+@pytest.mark.unit
+def test_world_mask_rejects_tagged_struct():
+    # A struct declares its own fields, so a core ndim on a field holding a struct is an error.
+    @flax.struct.dataclass
+    class TaggedStruct:
+        nested: Any = field(metadata={CORE_NDIM_KEY: 1})
+
+    with pytest.raises(ValueError, match="declares its own fields"):
+        world_mask(TaggedStruct(nested=TaggedStruct(nested=None)))
+
+
+@pytest.mark.unit
+def test_world_mask_infers_randomized_parameter():
+    # A parameter that is shared by default declares itself once it gains a world axis
+    sim = Sim(n_worlds=3)
+    assert not world_mask(sim.data).params.drag_matrix  # (3, 3), shared
+    drag = jnp.broadcast_to(sim.data.params.drag_matrix, (3, 1, 3, 3))
+    sim.data = sim.data.replace(params=sim.data.params.replace(drag_matrix=drag))
+    mask = world_mask(sim.data)
+    assert mask.params.drag_matrix  # (3, 1, 3, 3), now world-indexed
+    assert not mask.params.gravity_vec  # Siblings are unaffected