{"record":{"id":"7147c5262569fc4b","repo":"jax-ml/jax","slug":"kernel-input-j-in-hbm-but-does-not-have-trivial","errorCode":null,"errorMessage":"Kernel input {j} in HBM but does not have trivial BlockSpec.","messagePattern":"Kernel input (.+?) in HBM but does not have trivial BlockSpec\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":2368,"sourceCode":"          token = callback.io_callback(\n              # TODO(jburnim): Pass source_info from the pallas_call, in case this\n              # store is involved in a data race.\n              store,\n              TOKEN_SHAPE_DTYPE,\n              token,\n              device_id,\n              core_index,\n              TPU_MEMORY_SPACE_IDXS[memory_space],\n              input_ids[index],\n              (),\n              sliced_val,\n          )\n          return token\n\n        for j, var in enumerate(input_vars):\n          if input_var_memory_spaces[j] is _HBM:\n            if var.aval.shape != block_shapes[j]:\n              raise ValueError(\n                  f'Kernel input {j} in HBM but does not have trivial'\n                  ' BlockSpec.'\n              )\n            continue\n          assert len(cur_start_indices[j].shape) == 1\n          assert len(prev_start_indices[j].shape) == 1\n          token = jax.lax.cond(\n              (iteration_idx == initial_iteration_idx)\n              | jax.lax.reduce_or(\n                  cur_start_indices[j] != prev_start_indices[j], axes=(0,)\n              ),\n              functools.partial(\n                  _store_slice_to_kernel_input,\n                  j,\n                  var,\n                  input_var_memory_spaces[j],\n              ),\n              lambda t: t,","sourceCodeStart":2350,"sourceCodeEnd":2386,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L2350-L2386","documentation":"Inside the TPU interpret-mode grid loop, an input that lives in HBM must have a trivial BlockSpec: its block shape must equal the full input shape (the whole tensor is passed to the kernel each iteration). If the interpreter sees an HBM input whose kernel-visible block shape differs from the input shape, it raises this ValueError naming the input index.","triggerScenarios":"Passing an HBM input (e.g., via memory space annotations or mosaic params) with a BlockSpec that slices it into blocks, e.g. BlockSpec((BM,), ...) for an input of shape (N, BM).","commonSituations":"Marking an input as HBM while reusing the same blocked BlockSpec as VMEM inputs; kernels that stream blocks expecting HBM refs with per-iteration block views, which TPU interpret mode does not support.","solutions":["Give HBM inputs a BlockSpec with block shape equal to the full tensor shape (identity mapping), or None for scalars","Keep HBM inputs as copy sources (async_copy) rather than block-addressed arguments","Verify each input's memory space annotation matches its BlockSpec granularity"],"exampleFix":"# before\nhbm_spec = BlockSpec((BM, BN), lambda i, j: (i*BM, j*BN))  # blocked HBM input\n# after\nhbm_spec = None if scalar else BlockSpec((M, N), lambda i, j: (0, 0))  # whole tensor, trivial","handlingStrategy":"validation","validationCode":"for j, (x, spec) in enumerate(zip(inputs, in_specs)):\n    if memory_spaces[j] == 'hbm' and spec is not None:\n        assert tuple(spec.block_shape) == tuple(x.shape), f'input {j} in HBM needs trivial BlockSpec'","typeGuard":"def has_trivial_blockspec(spec, shape) -> bool:\n    return spec is None or tuple(spec.block_shape) == tuple(shape)","tryCatchPattern":"try:\n    interpret_run(kernel)\nexcept ValueError as e:\n    if 'does not have trivial BlockSpec' in str(e) and 'input' in str(e):\n        # replace input spec with whole-tensor trivial spec and retry\n        raise","preventionTips":["Give HBM args whole-tensor BlockSpecs or None","Stage HBM blocks via async_copy instead of blocked args","Review memory-space vs BlockSpec granularity per input"],"tags":["jax","pallas","tpu","hbm","blockspec","interpret-mode"],"backgroundTag":"invalid-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}