jax-ml/jax · error · ValueError

shard_map requires a non-empty mesh. Got {mesh}

Error message

shard_map requires a non-empty mesh. Got {mesh}

What it means

Even when a mesh object is supplied, it must be non-empty: it must contain at least one device/axis. An empty mesh (zero devices, e.g. Mesh(np.empty(0, dtype=object), axis_names=())) passes the isinstance check but fails here, since there is nothing to shard over.

Source

Thrown at jax/_src/shard_map.py:383

  if mesh is None:
    mesh = get_abstract_mesh()
    if mesh.empty:
      raise ValueError(
          "The context mesh cannot be empty. Use"
          " `jax.set_mesh(mesh)` to enter into a mesh context")
  else:
    ctx_mesh = get_abstract_mesh()
    if not ctx_mesh.empty and mesh.abstract_mesh != ctx_mesh:
      raise ValueError(
          f"The context mesh {ctx_mesh} should match the mesh passed to"
          f" shard_map {mesh}")

  if not isinstance(mesh, (Mesh, AbstractMesh)):
    raise TypeError("shard_map requires a `jax.sharding.Mesh` or a "
                    "`jax.sharding.AbstractMesh` instance for its "
                    f"second argument, but got {mesh} of type {type(mesh)}.")
  if mesh.empty:
    raise ValueError(f"shard_map requires a non-empty mesh. Got {mesh}")

  mesh_axis_names_wo_vmap = (
      frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names
  )

  if not isinstance(axis_names, (frozenset, set)):
    raise TypeError(
        "`axis_names` argument of shard_map should be of type `frozenset` or"
        f" `set`. Got type: {type(axis_names)}")
  if isinstance(axis_names, set):
    axis_names = frozenset(axis_names)
  if not axis_names:
    axis_names = mesh_axis_names_wo_vmap
  if not axis_names.issubset(mesh_axis_names_wo_vmap):
    raise ValueError(
        f"jax.shard_map requires axis_names={axis_names} to be a subset of "
        f"mesh.axis_names={mesh_axis_names_wo_vmap}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Check jax.devices() is non-empty before building the mesh and shard_map call
  2. Build the mesh from all available devices: Mesh(jax.devices(), axis_names=('i',))

Example fix

// before
mesh = jax.sharding.Mesh(np.empty((0,), dtype=object), axis_names=())
jax.shard_map(f, mesh=mesh, ...)(x)

// after
mesh = jax.sharding.Mesh(jax.devices(), axis_names=('i',))
jax.shard_map(f, mesh=mesh, ...)(x)
Defensive patterns

Strategy: validation

Validate before calling

assert len(jax.devices()) > 0, 'no devices available'
assert mesh.size > 0 and not mesh.empty, f'empty mesh: {mesh}'

Type guard

def is_nonempty_mesh(m) -> bool:
    return isinstance(m, jax.sharding.Mesh) and not m.empty

Prevention

When it happens

Trigger: Constructing a Mesh from an empty device array and passing it to shard_map, or programmatically building meshes from a device selection that returned no devices.

Common situations: Test harnesses or CPU-only environments where device enumeration unexpectedly yields nothing; scripts where devices are filtered by a predicate (e.g. GPU-only) that matches none; refactors leaving a placeholder empty mesh.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/ed8884b9a82d2a7c. Report an issue: GitHub.