{"record":{"id":"4e5b2ea1ca1319f9","repo":"jax-ml/jax","slug":"shape-polymorphism-for-pallas-does-not-support-dyn","errorCode":null,"errorMessage":"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.","messagePattern":"shape polymorphism for Pallas does not support dynamically-shaped blocks\\. Block spec for (.+?) has block_shape: (.+?)\\. If this is intentional, use the core\\.pallas_export_experimental\\(dynamic_shapes=True\\) context manager to enable dynamic shapes\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":638,"sourceCode":"      block_array_aval = array_aval.update(\n          shape=ref_block_shape, memory_space=jax_core.MemorySpace.Device,\n          sharding=ref_sharding)\n    elif isinstance(array_aval, state_types.AbstractLinVal):\n      if not isinstance(array_aval.inner_aval, jax_core.ShapedArray):\n        raise NotImplementedError  # TODO(mattjj,sharadmv)\n      block_array_aval = array_aval.inner_aval.update(shape=ref_block_shape)\n    else:\n      block_array_aval = array_aval.update(shape=ref_block_shape)\n    memory_space = self.memory_space\n    if memory_space is None:\n      memory_space = MemorySpace.DEFAULT\n    block_aval = state.AbstractRef(block_array_aval, memory_space)\n\n    if (\n        not jax_core.is_constant_shape(block_aval.shape)\n        and not dynamic_shapes_export_enabled()\n    ):\n      raise ValueError(\n          \"shape polymorphism for Pallas does not support \"\n          \"dynamically-shaped blocks. \"\n          f\"Block spec for {origin} has block_shape: {block_aval.shape}. \"\n          \"If this is intentional, use the \"\n          \"core.pallas_export_experimental(dynamic_shapes=True) \"\n          \"context manager to enable dynamic shapes.\"\n      )\n\n    fake_index_map_args, fake_index_map_kwargs = \\\n        index_map_tree.unflatten([False] * index_map_tree.num_leaves)\n    debug_info = api_util.debug_info(\n        \"pallas_call index_map\",\n        index_map_func,\n        fake_index_map_args,\n        fake_index_map_kwargs,\n    )\n    with tracing_grid_env(grid, vmapped_dims):\n      closed_jaxpr, out_avals = pe.trace_to_jaxpr(","sourceCodeStart":620,"sourceCodeEnd":656,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L620-L656","documentation":"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.","triggerScenarios":"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.","commonSituations":"Using jax2tf or aot_export with Pallas kernels on variable batch sizes; recent JAX versions where the guard was added for safety.","solutions":["Wrap the export in jax._src.pallas.core / jax.experimental.pallas export context: with pallas_export_experimental(dynamic_shapes=True): ...","Otherwise pin block shapes to concrete ints so blocks stay static","Fall back to non-polymorphic export (fixed shapes) if dynamic blocks aren't truly needed"],"exampleFix":"# before\nexported = jax.export(pallas_fn)(x)  # polymorphic shapes -> error\n# after\nfrom jax._src import pallas as pl_core\nwith pl_core.pallas_export_experimental(dynamic_shapes=True):\n    exported = jax.export(pallas_fn)(x)","handlingStrategy":"fallback","validationCode":"from jax._src import config as jcfg\nif exporting_with_polymorphic_shapes:\n    assert all(isinstance(d, int) for d in flatten(block_shape)), 'static blocks required'","typeGuard":null,"tryCatchPattern":"catch ValueError and retry export with dynamic_shapes=True or with fixed shapes","preventionTips":["Keep block dims concrete ints for export","Gate polymorphic Pallas export behind the experimental context manager"],"tags":["jax","pallas","shape-polymorphism","export","experimental"],"backgroundTag":"unsupported-feature-flag","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}