{"record":{"id":"870311ff8ba2bcc0","repo":"jax-ml/jax","slug":"jax-numpy-op-with-a-symbolic-number-of-sections","errorCode":null,"errorMessage":"jax.numpy.{op} with a symbolic number of sections is not supported","messagePattern":"jax\\.numpy\\.(.+?) with a symbolic number of sections is not supported","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":3158,"sourceCode":"    # and the result is clipped to [0, size] so that out-of-bound indices yield\n    # empty sections rather than negative sizes. Symbolic indices are left\n    # untouched, and clipping is skipped for symbolic sizes, since neither\n    # comparison is well-defined for them.\n    def _resolve(i_s):\n      i = core.concrete_dim_or_error(i_s, f\"in jax.numpy.{op} argument 1\")\n      if core.is_symbolic_dim(i):\n        return i\n      if i < 0:\n        i += size\n      if core.is_symbolic_dim(size):\n        return i\n      return np.clip(i, 0, size)\n    split_indices = np.asarray(\n        [0, *(_resolve(i_s) for i_s in indices_or_sections), size])\n    sizes = list(np.diff(split_indices))\n  else:\n    if core.is_symbolic_dim(indices_or_sections):\n      raise ValueError(f\"jax.numpy.{op} with a symbolic number of sections is \"\n                       \"not supported\")\n    num_sections: int = core.concrete_or_error(int, indices_or_sections,\n                                               f\"in jax.numpy.{op} argument 1\")\n    part_size, r = divmod(size, num_sections)\n    if r == 0:\n      sizes = [part_size] * num_sections\n    elif op == \"array_split\":\n      sizes = [(part_size + 1)] * r + [part_size] * (num_sections - r)\n    else:\n      raise ValueError(f\"array split does not result in an equal division: rest is {r}\")\n  sizes = [i if core.is_symbolic_dim(i) else np.int64(i)\n           for i in sizes]\n  return list(lax.split(ary, sizes, axis=axis))\n\n\n@export\ndef split(ary: ArrayLike, indices_or_sections: int | Sequence[int] | ArrayLike,\n          axis: int = 0) -> list[Array]:","sourceCodeStart":3140,"sourceCodeEnd":3176,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L3140-L3176","documentation":"jnp.split-family functions accept either an index list or an integer number of sections, but the integer form must be a concrete (statically known) value. A symbolic dimension (e.g. from jax.export or dynamic shapes) as the number of sections cannot be compiled.","triggerScenarios":"Calling jnp.split(x, n) where n is a symbolic dimension or tracer, e.g. splitting along an axis whose size or divisor comes from export-shaped inputs with dynamic dimensions.","commonSituations":"Using jax.export / dynamic shapes; passing a lazily-computed section count that is a DimVar rather than a Python int; splitting under transformations that lose concretization.","solutions":["Pass an explicit list of split indices instead of a section count: jnp.split(x, [i * step for i in range(1, n)]) with concrete ints","Make the number of sections a concrete Python int (compute it outside jit/export)","If size is symbolic too, compute split sizes as symbolic expressions via the index-list path"],"exampleFix":"// before\njnp.split(x, num_parts)  # num_parts is a symbolic dim\n// after\njnp.split(x, [k * (x.shape[0] // num_parts_) for k in range(1, num_parts_)])  # concrete ints","handlingStrategy":"fallback","validationCode":"if not isinstance(indices_or_sections, int):\n    indices_or_sections = int(indices_or_sections)  # or build index list","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pass explicit split index lists under dynamic shapes/jax.export","Compute section counts as concrete Python ints outside jit"],"tags":["jax","split","symbolic-dim","tracing"],"backgroundTag":"non-concrete-value-under-tracing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}