{"record":{"id":"342be4ca534bbefc","repo":"jax-ml/jax","slug":"only-slicing-with-static-indices-allowed-342be4","errorCode":null,"errorMessage":"Only slicing with static indices allowed","messagePattern":"Only slicing with static indices allowed","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":2064,"sourceCode":"    if isinstance(reg_type, ir.VectorType):\n      reg_shape = ir.VectorType(reg_type).shape\n      ty = ir.VectorType.get(reg_shape, elt)\n    else:\n      ty = elt\n\n    return self._pointwise(\n        lambda x: arith.bitcast(ty, x), output_is_signed=output_is_signed, restrict_bitwidth=False\n    )\n\n  def __getitem__(self, idx) -> FragmentedArray:\n    base_idx, slice_shape, is_squeezed = utils.parse_indices(idx, self.shape)\n    if isinstance(self.layout, WGSplatFragLayout):\n      shape = tuple(d for d, s in zip(slice_shape, is_squeezed) if not s)\n      return self.splat(self.registers.item(), shape, is_signed=self.is_signed)\n    if not isinstance(self.layout, TiledLayout):\n      raise NotImplementedError(\"Only arrays with tiled layouts can be sliced\")\n    if any(isinstance(idx, ir.Value) for idx in base_idx):\n      raise ValueError(\"Only slicing with static indices allowed\")\n    base_idx = cast(tuple[int, ...], base_idx)\n    base_tile_shape = self.layout.base_tile_shape\n    untiled_rank = len(self.shape) - len(base_tile_shape)\n    if any(is_squeezed[untiled_rank:]):\n      raise NotImplementedError(\n          \"Integer indexing not implemented for tiled dimensions (only slicing\"\n          \" allowed)\"\n      )\n    if untiled_rank:\n      base_tile_shape = (1,) * untiled_rank + base_tile_shape\n    if any(b % t for b, t in zip(base_idx, base_tile_shape, strict=True)):\n      raise ValueError(\n          \"Base indices of array slices must be aligned to the beginning of a\"\n          f\" tile. The array uses a tiling of {base_tile_shape}, but your base\"\n          f\" indices are {base_idx}. Consider using a different array layout.\"\n      )\n    if any(l % t for l, t in zip(slice_shape, base_tile_shape, strict=True)):\n      raise ValueError(","sourceCodeStart":2046,"sourceCodeEnd":2082,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L2046-L2082","documentation":"When slicing a FragmentedArray with a tiled layout, any index that is a dynamic ir.Value (produced at runtime) triggers ValueError('Only slicing with static indices allowed'). Mosaic slicing needs compile-time-constant base indices to compute register slices.","triggerScenarios":"fa[idx_value, :] where idx_value is an ir.Value (dynamic index computed inside the kernel), rather than a Python int or constant.","commonSituations":"Using loop-carried runtime offsets (e.g. from a dynamic loop index) as slice positions inside a Mosaic kernel.","solutions":["Use static Python-int indices for slicing","Compute the address/dynamic access via load/store with computed offsets instead of slicing","Hoist the slice out of dynamic control flow so indices are trace-time constants"],"exampleFix":"# before\nsub = fa[fi, :]  # fi is an ir.Value from loop induction\n# after\nsub = fa[0, :]  # static index; or gather via memory ops with dynamic offset\nptr = base_ptr + fi * stride  # dynamic addressing at memory level","handlingStrategy":"validation","validationCode":"assert all(not isinstance(i, ir.Value) for i in base_idx), 'static indices only'","typeGuard":"def all_static(idx) -> bool:\n    import ir\n    flat = idx if isinstance(idx, tuple) else (idx,)\n    return all(not isinstance(i, ir.Value) for i in flat)","tryCatchPattern":null,"preventionTips":["Use Python int constants for slice positions","Handle dynamic offsets via memory addressing, not fragment slicing"],"tags":["mosaic","gpu","slicing","static-indices"],"backgroundTag":"dynamic-index-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}