jax-ml/jax · error · ValueError

Sharding has to be concrete when layout is of type {type(lay

Error message

Sharding has to be concrete when layout is of type {type(layout)}. Please pass a `jax.sharding.NamedSharding` or `jax.sharding.SingleDeviceSharding` to the sharding argument. Got sharding {sharding}

What it means

When constructing a Layout with an explicit jax Layout object, a concrete sharding (NamedSharding or SingleDeviceSharding) must also be supplied. Layout without sharding is rejected.

Source

Thrown at jax/_src/layout.py:138

  def check_compatible_aval(self, aval_shape: Shape):
    if len(self.major_to_minor) != len(aval_shape):
      raise ValueError(
          f'Length of major_to_minor and the rank of the value should match.'
          f' Got major_to_minor={self.major_to_minor} and shape={aval_shape}')


LayoutOptions = Layout | None | AutoLayoutSingleton
ShardingOptions = Sharding | None


class Format:
  __slots__ = ['layout', 'sharding']

  def __init__(self, layout: LayoutOptions = None,
               sharding: ShardingOptions = None):
    # If layout is concrete and sharding is not, error.
    if isinstance(layout, Layout) and sharding is None:
      raise ValueError(
          'Sharding has to be concrete when layout is of type'
          f' {type(layout)}. Please pass a'
          ' `jax.sharding.NamedSharding` or'
          ' `jax.sharding.SingleDeviceSharding` to the sharding argument. Got'
          f' sharding {sharding}'
      )
    if not isinstance(
        layout, (Layout, type(None), AutoLayoutSingleton)):
      raise TypeError(
          'Invalid value received for the layout argument.'
          ' Expected values are `None`, `Layout.AUTO` or an'
          f' instance of `Layout`. Got {layout} of'
          f' type {type(layout)}'
      )
    if not isinstance(sharding, (Sharding, type(None))):
      raise TypeError(
          'Invalid value received for the sharding argument. Expected values'
          ' are `None` or an instance of `jax.Sharding`. Got'

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass a concrete sharding: jax.sharding.NamedSharding(mesh, P('x')) or jax.sharding.SingleDeviceSharding(jax.devices()[0])
  2. Or omit the explicit Layout and let JAX infer layout from the sharding (AUTO)

Example fix

# before
Layout(Layout(major_to_minor=(1,0)), sharding=None)
# after
import jax
Layout(Layout(major_to_minor=(1,0)),
       sharding=jax.sharding.SingleDeviceSharding(jax.devices()[0]))
Defensive patterns

Strategy: validation

Validate before calling

import jax
from jax._src.layout import Layout
if isinstance(layout, Layout):
    assert sharding is not None and isinstance(sharding, jax.sharding.Sharding)

Type guard

import jax
from jax._src.layout import Layout
def has_concrete_sharding(layout, sharding):
    return not isinstance(layout, Layout) or isinstance(sharding, jax.sharding.Sharding)

Prevention

When it happens

Trigger: Layout(layout=Layout(major_to_minor=(1,0)), sharding=None) in APIs taking layout/sharding pairs (e.g. jax device placement or array creation options).

Common situations: Specifying memory layouts on multi-device JAX arrays without stating how the array is sharded.

Related errors


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