{"record":{"id":"a7b8100a1ab087e6","repo":"jax-ml/jax","slug":"kernel-output-j-in-hbm-but-does-not-have-trivial","errorCode":null,"errorMessage":"Kernel output {j} in HBM but does not have trivial BlockSpec.","messagePattern":"Kernel output (.+?) 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":2432,"sourceCode":"              functools.partial(store, output_name=output_names[index]),\n              TOKEN_SHAPE_DTYPE,\n              token,\n              device_id,\n              core_index,\n              TPU_MEMORY_SPACE_IDXS[mosaic_core.MemorySpace.HBM],\n              output_buffer_ids[index],\n              (transform,),\n              kernel_output_val,\n              cur_block_indices[num_inputs + index],\n              grid_point,\n          )\n          return token\n\n        output_slices : list[Any] = []\n        for j, var in enumerate(output_vars):\n          if output_var_memory_spaces[j] is _HBM:\n            if var.aval.shape != block_shapes[num_inputs + j]:\n              raise ValueError(\n                  f'Kernel output {j} in HBM but does not have trivial'\n                  ' BlockSpec.'\n              )\n            output_slices.append(None)\n            continue\n          assert len(cur_start_indices[num_inputs + j].shape) == 1\n          assert len(next_start_indices[num_inputs + j].shape) == 1\n          transform = indexing.NDIndexer(\n              indices=tuple(\n                  indexing.ds(st, sz) if not iid else st\n                  for st, sz, iid in zip(\n                      cur_start_indices[num_inputs + j],\n                      block_shapes[num_inputs + j],\n                      is_squeeze_dim[num_inputs + j],\n                  )\n              ),\n              shape=output_vals[j].shape,\n              int_indexer_shape=(),","sourceCodeStart":2414,"sourceCodeEnd":2450,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L2414-L2450","documentation":"Symmetric to the input case: in TPU interpret mode, an output residing in HBM must have a trivial BlockSpec whose block shape equals the full output shape. If the kernel's block shape for that output differs from the output tensor shape, this ValueError is raised naming the output index.","triggerScenarios":"Declaring an HBM output with a blocked BlockSpec (block shape smaller than out shape) in a kernel run under TPU interpret mode.","commonSituations":"Streaming-style kernels writing blocks directly to HBM outputs; reusing one grid/BlockSpec tuple for inputs and outputs where only inputs should be blocked; GPU Pallas kernels ported to TPU interpretation.","solutions":["Use a trivial BlockSpec (full shape, zero start indices) for HBM outputs","Write results via async_copy from a VMEM buffer to the HBM output instead of blocked stores","Split outputs so blocked ones are VMEM-intermediates copied back manually"],"exampleFix":"# before\nout_spec = BlockSpec((BM, BN), lambda i, j: (i*BM, j*BN))  # HBM output blocked\n# after\nout_spec = BlockSpec((M, N), lambda i, j: (0, 0))  # trivial for HBM output","handlingStrategy":"validation","validationCode":"for j, (o, spec) in enumerate(zip(out_shapes, out_specs)):\n    if output_memory_spaces[j] == 'hbm' and spec is not None:\n        assert tuple(spec.block_shape) == tuple(o.shape), f'output {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 'output' in str(e):\n        # use full-shape BlockSpec for that output and retry\n        raise","preventionTips":["Use trivial BlockSpecs for HBM outputs","Copy results from VMEM with async_copy instead of blocked stores","Keep blocked specs only for non-HBM buffers"],"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"}