jax-ml/jax · error · ValueError

mode='slice' is not valid for polymorphic shapes.

Error message

mode='slice' is not valid for polymorphic shapes.

What it means

to_static_slice refuses to produce a static slice when any dimension of the array's shape is symbolic (polymorphic, e.g. from jax2ts or shape polymorphism with axis names like 'n'). Static slicing requires concrete sizes, so mode='slice' paths bail out with this ValueError.

Source

Thrown at jax/_src/numpy/indexing.py:382

      else:
        new_indices.append(idx)
    return NDIndexer(indices=new_indices, shape=self.shape)

  def to_static_slice(
      self, *,
      arr_is_sharded: bool = False,
      normalize_indices: bool = True,
      mode: str | slicing.GatherScatterMode | None) -> _StaticSliceIndexer:
    """Convert to StaticSliceIndexer data structure.

    If this is not possible, raise a ValueError, TypeError, or IndexError.
    """
    if mode is None:
      parsed_mode = slicing.GatherScatterMode.PROMISE_IN_BOUNDS
    else:
      parsed_mode = slicing.GatherScatterMode.from_any(mode)
    if any(core.is_symbolic_dim(s) for s in self.shape):
      raise ValueError("mode='slice' is not valid for polymorphic shapes.")

    if parsed_mode not in [
        slicing.GatherScatterMode.PROMISE_IN_BOUNDS, slicing.GatherScatterMode.CLIP]:
      raise ValueError("static_slice requires mode='promise_in_bounds' or mode='clip'")

    # Validation of the unmodified user indices.
    if parsed_mode == slicing.GatherScatterMode.PROMISE_IN_BOUNDS:
      self.validate_static_indices(normalize_indices=normalize_indices)
    self.validate_slices()

    # For sharded inputs, indexing (like x[0]) and partial slices (like x[:2] as
    # opposed to x[:]) lead to incorrect sharding semantics when computed via slice.
    # TODO(yashkatariya): fix slice with sharding
    if arr_is_sharded and self.has_partial_slices():
      raise ValueError("static_slice with partial slices does not support nontrivial array sharding.")

    for position, pidx in enumerate(self.indices):
      if pidx.typ in [IndexType.INTEGER, IndexType.ELLIPSIS, IndexType.SLICE, IndexType.NONE]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Avoid slicing the polymorphic axis; use full-axis ops or lax.dynamic_slice with static size where supported
2. Mark that dimension as static (concrete int) if slicing is essential
3. Reorganize computation so the sliced output size remains static

When it happens

Trigger: Using shape-polymorphic jit (e.g. jax.jit(f, axis_arg/ shape annotations like 'n') or jax2ts export) and indexing/slicing x within the traced function via the static-slice path.

Common situations: JAX shape-polymorphism / TorchExport workflows where a dynamic batch dim ('n') is declared and code slices that axis.

Related errors

• multi-platform lowering for buffer_callback
• No VJP is available
• Duplicate serialization registration for type `{nodetype}`.
• Duplicate serialization registration for serialized_name `{s
• serializing DisabledSafetyCheck: {check}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/a259bf9aba76fd1a. Report an issue: GitHub.