{"record":{"id":"6ef285564d1023c5","repo":"jax-ml/jax","slug":"get-not-supported-yet-for-block-shape-b","errorCode":null,"errorMessage":"get not supported yet for block shape {b}","messagePattern":"get not supported yet for block shape (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":1697,"sourceCode":"    raise NotImplementedError('get not supported yet')\n  if not indexers:\n    indexer = indexing.NDIndexer.make_trivial_indexer(ref_aval.shape)\n    indexer_aval = indexer\n  else:\n    indexer = indexers[0]\n    indexer_aval = indexers_avals[0]\n  block_indexer = []\n\n  def _slice(i, b):\n    match b:\n      case int():\n        return indexing.ds(i * b, b)\n      case pallas_core.Blocked(bs):\n        return indexing.ds(i * bs, bs)\n      case pallas_core.Squeezed() | None:\n        return i\n      case _:\n        raise NotImplementedError(f'get not supported yet for block shape {b}')\n\n  if (\n      ref_block_spec is pallas_core.no_block_spec\n      or ref_block_spec.block_shape is None\n  ):\n    # Short-circuit if the ref is not blocked.\n    return state_primitives.get_p.bind(ref, *idx, tree=tree)\n  block_idx_iter = iter(ctx.get_out_block_indices()[0])\n  for idx_aval, size, idx, bd in zip(\n      indexer_aval.indices,\n      ref_aval.shape,\n      indexer.indices,\n      ref_block_spec.block_shape,\n      strict=True,\n  ):\n    if not isinstance(idx_aval, indexing.Slice):\n      assert hasattr(idx_aval, 'shape') and not idx_aval.shape, idx_aval\n      assert bd is None or isinstance(bd, pallas_core.Squeezed)","sourceCodeStart":1679,"sourceCodeEnd":1715,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L1679-L1715","documentation":"Raised by the _slice helper inside _get_eval_rule when the ref's block_shape entry for a dimension is not int, pallas_core.Blocked, pallas_core.Squeezed, or None — typically pallas_core.Element. The eval rule cannot compute the per-block index (i * b) for such block descriptors when materializing a get.","triggerScenarios":"ref.get / ref[...] on a Ref blocked with pallas_core.Element (or any exotic block descriptor) inside a pallas kernel evaluated by the fuser.","commonSituations":"Using Element block specs for fine-grained layouts then reading the ref; mixing Element-blocked outputs with gets in fused kernels.","solutions":["Change the BlockSpec to use int/Blocked/Squeezed block descriptors on dims that get read","Avoid get on Element-blocked refs; read via an int-blocked alias buffer"],"exampleFix":"# before\nBlockSpec(block_shape=(pallas_core.Element(1), 64), ...)\nv = ref[:]\n\n# after\nBlockSpec(block_shape=(1, 64), ...)\nv = ref[:]","handlingStrategy":"type-guard","validationCode":"import jax._src.pallas.pallas_core as pc\nassert all(b is None or isinstance(b, (int, pc.Blocked, pc.Squeezed)) for b in block_shape), f'unsupported block type for get: {block_shape}'","typeGuard":"def is_gettable_block(b) -> bool:\n    import jax._src.pallas.pallas_core as pc\n    return b is None or isinstance(b, (int, pc.Blocked, pc.Squeezed))","tryCatchPattern":null,"preventionTips":["Keep Element blocks out of read paths","Standardize on int block sizes for refs consumed via get"],"tags":["jax","pallas","element-indexing","get","not-implemented"],"backgroundTag":"jax-pallas-unsupported-indexing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}