{"record":{"id":"75068a333d32b6d5","repo":"jax-ml/jax","slug":"unsupported-memory-space-x-memory-space","errorCode":null,"errorMessage":"Unsupported memory space: {x.memory_space}","messagePattern":"Unsupported memory space: (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":1896,"sourceCode":"          subkey, coordindates_along_dim\n      )\n      grid_point_coordinates.append(coordindates_along_dim)\n    else:\n      grid_point_coordinates.append(jnp.array((), dtype=jnp.int32))\n\n  return tuple(grid_point_coordinates)\n\n# TODO(sharadmv, jburnim): add support for memory space constraints\nremove_memory_space_p = jax_core.Primitive('remove_memory_space')\n\n@remove_memory_space_p.def_abstract_eval\ndef _remove_memory_space_abstract_eval(x):\n  if (isinstance(x, jax_core.ShapedArray) and\n      not isinstance(x.memory_space, jax_core.MemorySpace)):\n    if (x.memory_space is None or x.memory_space is pallas_core.MemorySpace.ANY\n        or x.memory_space is mosaic_core.MemorySpace.HBM):\n      return jax_core.ShapedArray(x.shape, x.dtype)\n    raise NotImplementedError(f'Unsupported memory space: {x.memory_space}')\n  return x\n\n@remove_memory_space_p.def_impl\ndef _remove_memory_space_impl(x):\n  return x\n\ndef _remove_memory_space_lowering(_, x):\n  return [x]\nmlir.register_lowering(remove_memory_space_p, _remove_memory_space_lowering)\n\n\ndef _get_grid_point(\n    loop_indices: tuple[Array, ...],\n    grid_point_coordinates: _GridPointCoordinatesPerDim,\n) -> Array:\n  \"\"\"Indexes each entry in `grid_point_coordinates` with the corresponding entry in `loop_indices`.\n\n  If an entry in `grid_point_coordinates` is an empty array, the corresponding","sourceCodeStart":1878,"sourceCodeEnd":1914,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L1878-L1914","documentation":"The remove_memory_space abstract eval strips memory-space annotations from arrays, but only supports memory_space of None, ANY, or HBM (mosaic). Any other memory space annotation raises NotImplementedError, since the interpreter cannot erase an unrecognized memory space.","triggerScenarios":"Passing an array whose memory_space attribute is an unsupported enum value (e.g., a mosaic VMEM/SMEM annotated value) through the remove_memory_space_p primitive during interpretation setup.","commonSituations":"Feeding pre-annotated buffers (VMEM refs) directly as pallas_call inputs/outputs instead of plain HBM arrays; version skew where new memory-space enums exist in mosaic but not the interpreter path.","solutions":["Pass plain jax arrays (no memory space annotation) as pallas_call inputs/outputs","Allocate VMEM/SMEM inside the kernel via run_scoped, not as external arguments","Update JAX/mosaic versions so supported memory spaces agree"],"exampleFix":"# before\nouts = pallas_call(kernel, out_shapes, vmem_annotated_inputs)  # annotated arrays\n# after\nouts = pallas_call(kernel, out_shapes, plain_jnp_arrays)","handlingStrategy":"type-guard","validationCode":"for a in jax.tree_util.tree_leaves(args):\n    ms = getattr(a, 'memory_space', None)\n    assert ms is None or 'ANY' in str(ms) or 'HBM' in str(ms), f'unsupported memory_space {ms}'","typeGuard":"def has_supported_memory_space(a) -> bool:\n    ms = getattr(a, 'memory_space', None)\n    return ms is None or 'ANY' in str(ms) or 'HBM' in str(ms)","tryCatchPattern":"try:\n    interpret_run(kernel, args)\nexcept NotImplementedError as e:\n    if 'Unsupported memory space' in str(e):\n        args = jax.tree.map(lambda x: jnp.asarray(x) if hasattr(x, 'memory_space') else x, args)\n        interpret_run(kernel, args)","preventionTips":["Pass plain arrays as pallas_call inputs/outputs","Allocate VMEM/SMEM inside kernels","Keep jax and mosaic versions in sync"],"tags":["jax","pallas","memory-space","not-implemented"],"backgroundTag":"unsupported-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}