jax-ml/jax · error · ValueError

shape polymorphism for Pallas does not support dynamically-s

Error message

shape polymorphism for Pallas does not support dynamically-shaped blocks. Block spec for {origin} has block_shape: {block_aval.shape}. If this is intentional, use the core.pallas_export_experimental(dynamic_shapes=True) context manager to enable dynamic shapes.

What it means

Pallas export normally requires statically-known block shapes. If the block spec's block aval contains dynamic (polymorphic) dimensions — e.g. from jax.export with symbolic shapes — and the experimental dynamic-shapes export flag is off, to_block_mapping refuses to proceed and points to the core.pallas_export_experimental(dynamic_shapes=True) context manager.

Source

Thrown at jax/_src/pallas/core.py:638

      block_array_aval = array_aval.update(
          shape=ref_block_shape, memory_space=jax_core.MemorySpace.Device,
          sharding=ref_sharding)
    elif isinstance(array_aval, state_types.AbstractLinVal):
      if not isinstance(array_aval.inner_aval, jax_core.ShapedArray):
        raise NotImplementedError  # TODO(mattjj,sharadmv)
      block_array_aval = array_aval.inner_aval.update(shape=ref_block_shape)
    else:
      block_array_aval = array_aval.update(shape=ref_block_shape)
    memory_space = self.memory_space
    if memory_space is None:
      memory_space = MemorySpace.DEFAULT
    block_aval = state.AbstractRef(block_array_aval, memory_space)

    if (
        not jax_core.is_constant_shape(block_aval.shape)
        and not dynamic_shapes_export_enabled()
    ):
      raise ValueError(
          "shape polymorphism for Pallas does not support "
          "dynamically-shaped blocks. "
          f"Block spec for {origin} has block_shape: {block_aval.shape}. "
          "If this is intentional, use the "
          "core.pallas_export_experimental(dynamic_shapes=True) "
          "context manager to enable dynamic shapes."
      )

    fake_index_map_args, fake_index_map_kwargs = \
        index_map_tree.unflatten([False] * index_map_tree.num_leaves)
    debug_info = api_util.debug_info(
        "pallas_call index_map",
        index_map_func,
        fake_index_map_args,
        fake_index_map_kwargs,
    )
    with tracing_grid_env(grid, vmapped_dims):
      closed_jaxpr, out_avals = pe.trace_to_jaxpr(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the export in jax._src.pallas.core / jax.experimental.pallas export context: with pallas_export_experimental(dynamic_shapes=True): ...
2. Otherwise pin block shapes to concrete ints so blocks stay static
3. Fall back to non-polymorphic export (fixed shapes) if dynamic blocks aren't truly needed

Example fix

# before
exported = jax.export(pallas_fn)(x)  # polymorphic shapes -> error
# after
from jax._src import pallas as pl_core
with pl_core.pallas_export_experimental(dynamic_shapes=True):
    exported = jax.export(pallas_fn)(x)

from jax._src import config as jcfg
if exporting_with_polymorphic_shapes:
    assert all(isinstance(d, int) for d in flatten(block_shape)), 'static blocks required'

catch ValueError and retry export with dynamic_shapes=True or with fixed shapes

When it happens

Trigger: Exporting/jitting a pallas_call under shape polymorphism (jax.export or symbolic dimension formulas) so that block dims become non-constant, without enabling dynamic_shapes export.

Common situations: Using jax2tf or aot_export with Pallas kernels on variable batch sizes; recent JAX versions where the guard was added for safety.

Related errors

• Encountered dimension variable '{self.var}' that is not appe
• Cannot divide {self} by {divisor}.
• Invalid mixing of symbolic scopes {when}.\nExpected {self_de
• Encountered unexpected shape dimension {d}
• multi-platform lowering for buffer_callback

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