{"record":{"id":"4f255939f99023ea","repo":"jax-ml/jax","slug":"dimension-must-be-smaller-than-the-rank-of-the-a","errorCode":null,"errorMessage":"`dimension` must be smaller than the rank of the array.","messagePattern":"`dimension` must be smaller than the rank of the array\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":1162,"sourceCode":"\n    return cls(\n        _registers=np.full(layout.registers_shape(shape), value, dtype=object),\n        _layout=layout,\n        _is_signed=is_signed,\n    )\n\n  @staticmethod\n  def broadcasted_iota(\n      dtype: ir.Type,\n      shape: tuple[int, ...],\n      dimension: int,\n      layout: FragmentedLayout | None = None,\n      *,\n      is_signed: bool | None = None,\n  ) -> FragmentedArray:\n    \"\"\"Creates a broadcasted iota array along the specified dimension.\"\"\"\n    if dimension >= len(shape):\n      raise ValueError(\n          \"`dimension` must be smaller than the rank of the array.\"\n      )\n\n    def cast(idx: ir.Value) -> ir.Value:\n      if isinstance(dtype, ir.FloatType):\n        i32 = ir.IntegerType.get_signless(32)\n        return arith.uitofp(dtype, arith.index_cast(i32, idx))\n      return arith.index_cast(dtype, idx)\n\n    return mgpu.FragmentedArray.splat(\n        llvm.mlir_undef(dtype),\n        shape,\n        layout,\n        is_signed=is_signed,\n    ).foreach(\n        lambda _, idx: cast(idx[dimension]),\n        create_array=True,\n        is_signed=is_signed,","sourceCodeStart":1144,"sourceCodeEnd":1180,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L1144-L1180","documentation":"FragmentedArray.broadcasted_iota creates an iota (index sequence) distributed along one dimension of a register-level array on GPU. The `dimension` argument selects which axis the iota counts along, so it must index into the array's shape. If dimension >= len(shape) the axis doesn't exist and the lowering cannot proceed.","triggerScenarios":"Calling FragmentedArray.broadcasted_iota(dtype, shape, dimension=...) with a dimension index greater than or equal to rank(shape), e.g. shape=(8,) with dimension=1, or passing a dimension for a scalar shape ().","commonSituations":"Dynamically computing the iota dimension from loop variables or user-supplied rank in a Mosaic GPU kernel; off-by-one errors when dimension is derived from len(shape); refactoring a kernel from 2-D to 1-D tiles without updating the hardcoded dimension.","solutions":["Check that 0 <= dimension < len(shape) before calling broadcasted_iota and fix the off-by-one","If you need a higher-rank iota, pass a shape with rank > dimension (e.g. append a size-1 axis)","Log shape and dimension at kernel-build time to catch rank mismatches early"],"exampleFix":"# before\nfa = FragmentedArray.broadcasted_iota(i32, (16,), dimension=1)\n# after\nfa = FragmentedArray.broadcasted_iota(i32, (16, 1), dimension=1)","handlingStrategy":"validation","validationCode":"assert 0 <= dimension < len(shape), f\"dimension {dimension} out of range for shape {shape}\"","typeGuard":"def valid_iota_dim(shape: tuple[int, ...], dim: int) -> bool:\n    return isinstance(dim, int) and 0 <= dim < len(shape)","tryCatchPattern":null,"preventionTips":["Compute iota dimensions as len(shape)-relative expressions, never absolute constants","Add shape/dimension asserts in kernel-building helpers"],"tags":["jax","mosaic-gpu","iota","dimension-validation","shape-mismatch"],"backgroundTag":"axis-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}