{"record":{"id":"ffbd53e96bc49eff","repo":"jax-ml/jax","slug":"unsupported-block-dim-type-type-b-ffbd53","errorCode":null,"errorMessage":"Unsupported block dim type: {type(b)}","messagePattern":"Unsupported block dim type: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":1788,"sourceCode":"  jaxpr = block_mapping.index_map_jaxpr\n  token, block_indices = _interpret_jaxpr(\n      jaxpr,\n      *jaxpr.consts,\n      *loop_idx,\n      *args,\n      ctx=ctx,\n      token=token,\n  )\n  def _get_start_index(i, b):\n    match b:\n      case pallas_core.Squeezed():\n        return i\n      case pallas_core.Element():\n        return i\n      case pallas_core.Blocked():\n        return i * b.block_size\n      case _:\n        raise ValueError(f\"Unsupported block dim type: {type(b)}\")\n  ret = jnp.array(\n      tuple(\n          _get_start_index(i, b)\n          for i, b in zip(block_indices, block_mapping.block_shape)\n      ),\n      dtype=jnp.int32,\n  )\n  return token, block_indices, ret\n\n\ndef _get_parallel_dim_semantics(\n    mosaic_params: mosaic_core.CompilerParams, num_dimensions_in_grid: int,\n) -> tuple[bool, ...]:\n  \"\"\"Returns a tuple indicating which grid dimensions have parallel semantics.\n\n  Args:\n    mosaic_params: The compiler params for the Mosaic TPU backend.\n    num_dimensions_in_grid: The number of dimensions in the grid.","sourceCodeStart":1770,"sourceCodeEnd":1806,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L1770-L1806","documentation":"When the interpreter computes block start indices from a BlockSpec's block mapping, each dimension's block type must be one of pallas_core.Element, Blocked, or an int. Any other type in the BlockSpec's block_shape raises this ValueError naming the offending type.","triggerScenarios":"Constructing a BlockSpec whose block_shape entries are unexpected types (e.g., None, a string, a custom class, or numpy scalars in older versions) so the match statement in _get_start_index falls through to the error case.","commonSituations":"Passing None for broadcast/whole-tensor dims instead of pallas_core.Element; version mismatches between jax and jaxlib/pallas where BlockSpec internals changed; dynamically built block specs with heterogeneous entries.","solutions":["Ensure every block_shape entry is an int, pallas_core.Element(), or pallas_core.Blocked(...)","Use None (no BlockSpec) for scalar/whole-tensor arguments rather than odd block types","Align JAX/jaxlib versions so pallas_core types match the interpreter"],"exampleFix":"# before\nspec = BlockSpec((None, 128), lambda i: (0, i*128))\n# after\nfrom jax._src.pallas import pallas_core\nspec = BlockSpec((pallas_core.Element(), 128), lambda i: (0, i*128))","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas import pallas_core\nvalid = (int, pallas_core.Element, pallas_core.Blocked)\nassert all(isinstance(b, valid) for b in spec.block_shape), f'bad block types: {[type(b) for b in spec.block_shape]}'","typeGuard":"def is_valid_block_shape(bs) -> bool:\n    from jax._src.pallas import pallas_core\n    return all(isinstance(b, (int, pallas_core.Element, pallas_core.Blocked)) for b in bs)","tryCatchPattern":"try:\n    interpret_run(kernel)\nexcept ValueError as e:\n    if 'Unsupported block dim type' in str(e):\n        # normalize block_shape entries to int/Element/Blocked and retry\n        raise","preventionTips":["Build BlockSpecs only from documented block types","Use None BlockSpec for scalars/whole tensors","Pin jax/jaxlib versions to matching releases"],"tags":["jax","pallas","blockspec","type-error"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}