jax-ml/jax · error · ValueError

jnp.insert(): index array must be integer typed; got {obj}

Error message

jnp.insert(): index array must be integer typed; got {obj}

What it means

Raised by jnp.insert when the insertion positions array has a non-integer dtype (e.g. float) and is non-empty (or is a JAX Array). NumPy deprecates boolean/float obj for insert, so JAX rejects non-integer index arrays outright, only casting empty non-Array inputs to int.

Source

Thrown at jax/_src/numpy/lax_numpy.py:7802

  if axis is None:
    a = ravel(a)
    axis = 0
  axis = core.concrete_or_error(None, axis, "axis argument of jnp.insert()")
  axis = _canonicalize_axis(axis, a.ndim)
  if isinstance(obj, slice):
    indices = arange(*obj.indices(a.shape[axis]))
  else:
    indices = asarray(obj)

  if indices.ndim > 1:
    raise ValueError("jnp.insert(): obj must be a slice, a one-dimensional "
                     f"array, or a scalar; got {obj}")
  if not np.issubdtype(indices.dtype, np.integer):
    if indices.size == 0 and not isinstance(obj, Array):
      indices = indices.astype(int)
    else:
      # Note: np.insert allows boolean inputs but the behavior is deprecated.
      raise ValueError("jnp.insert(): index array must be "
                       f"integer typed; got {obj}")
  values_arr = array(values_arr, ndmin=a.ndim, dtype=a.dtype, copy=False)

  if indices.size == 1:
    index = ravel(indices)[0]
    if indices.ndim == 0:
      values_arr = moveaxis(values_arr, 0, axis)
    indices = array_creation.full(values_arr.shape[axis], index)
  n_input = a.shape[axis]
  n_insert = broadcast_shapes(indices.shape, (values_arr.shape[axis],))[0]
  out_shape = list(a.shape)
  out_shape[axis] += n_insert
  out = array_creation.zeros_like(a, shape=tuple(out_shape))

  indices = where(indices < 0, indices + n_input, indices)
  indices = clip(indices, 0, n_input)

  values_ind = indices.at[argsort(indices)].add(arange(n_insert, dtype=indices.dtype))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Cast positions to int: jnp.insert(a, obj.astype(int), values)
  2. Fix position computation to use integer arithmetic (//, operator.index)
  3. Convert boolean masks to positions with jnp.where(mask)[0] before inserting

Example fix

// before
jnp.insert(a, len(a)/2, 99)  # float index
// after
jnp.insert(a, len(a)//2, 99)  # or int(len(a)/2)
Defensive patterns

Strategy: type-guard

Validate before calling

obj = jnp.asarray(obj)
if not jnp.issubdtype(obj.dtype, jnp.integer) and obj.size:
    obj = obj.astype(jnp.int32)

Type guard

def is_integer_insert_positions(obj):
    o = jnp.asarray(obj)
    return o.size == 0 or jnp.issubdtype(o.dtype, jnp.integer)

Prevention

When it happens

Trigger: jnp.insert(a, jnp.array([0.5]), values); insertion positions computed as floats (division instead of integer division); boolean obj arrays, which NumPy allowed but deprecated.

Common situations: Positions like n//2 accidentally written n/2; positions from np.linspace without astype(int); porting old NumPy code using boolean masks with insert.

Related errors


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