jax-ml/jax · error · ValueError

The layout of ShapedArray should not be `AutoLayout` when la

Error message

The layout of ShapedArray should not be `AutoLayout` when layout mode is {cur_layout_mode}

What it means

get_layout rejects AutoLayout (the placeholder 'deferred' layout) when the global layout mode is not AUTO. In explicit/deferred layout modes, an aval must carry a concrete layout, not the Auto sentinel.

Source

Thrown at jax/_src/core.py:2449

                       unreduced_kind=kind)

  @property
  def empty(self):
    return self is empty_mat

  def invarying(self, mesh) -> frozenset:
    return frozenset(mesh.manual_axes) - (
        self.varying | self.unreduced | self.reduced)

  @property
  def vur(self) -> frozenset:
    return self.varying | self.unreduced | self.reduced

def get_layout(layout):
  cur_layout_mode = get_layout_mode()
  if (cur_layout_mode is not LayoutMode.AUTO and
      isinstance(layout, AutoLayoutSingleton)):
    raise ValueError(
        "The layout of ShapedArray should not be `AutoLayout` when layout mode"
        f" is {cur_layout_mode}")
  return layout


empty_mat = ManualAxisType()

@functools.cache
def _empty_sharding(ndim):
  return NamedSharding(mesh_lib.empty_abstract_mesh, P(*[None] * ndim))


@immutable
class ShapedArray(AbstractValue):
  # inherits slots from parent
  __slots__ = ['shape', 'dtype', 'weak_type', 'sharding', 'manual_axis_type',
               'memory_space', 'layout', '_stripped_weak_type', '__weakref__']
  array_abstraction_level = 2

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Supply a concrete layout instead of AutoLayout when layout mode is explicit
2. Avoid flipping global layout mode around code that creates default avals
3. Wrap mode-dependent construction so AUTO mode is only used where AutoLayout is allowed

Example fix

// before
aval = ShapedArray._create(shape, dtype, ..., layout=AutoLayout())  # under EXPLICIT mode

// after
aval = ShapedArray._create(shape, dtype, ..., layout=concrete_layout)

from jax._src.core import get_layout_mode, AutoLayoutSingleton
if get_layout_mode() is not LayoutMode.AUTO and isinstance(layout, AutoLayoutSingleton):
    layout = concrete_layout

When it happens

Trigger: Constructing avals with layout=AutoLayout while jax.config or the layout mode context sets LayoutMode.EXPLICIT (e.g. via jax.layout_mode or experimental layout APIs); usually in custom primitives or tests toggling modes.

Common situations: Enabling experimental layout inference/mode flags; code that worked under default AUTO mode then running under a different mode; mixing explicit-layout APIs with default avals.

Related errors

• Explicitly requested dtype {}{} is not available. To enable
• Expected mode of type `LayoutMode`. Got type: {type(mode)}
• Length of major_to_minor and the rank of the value should ma
• Sharding has to be concrete when layout is of type {type(lay
• Invalid value received for the layout argument. Expected val

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