jax-ml/jax · error · TypeError

offset must be an integer, got {offset!r}

Error message

offset must be an integer, got {offset!r}

What it means

jax.lax._tri (used by jnp.tri and triangular-mask helpers) requires the diagonal offset to be an integer type. If offset's dtype is not a subtype of np.integer (e.g. a float like 1.0 or a tracer of float dtype), it raises TypeError.

Source

Thrown at jax/_src/lax/lax.py:3709

def _delta(dtype: DTypeLike, shape: Shape, axes: Sequence[int]) -> Array:
  """This utility function exists for creating Kronecker delta arrays."""
  axes = map(int, axes)
  dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "delta")
  base_shape = tuple(np.take(shape, axes))
  iotas = [broadcasted_iota(np.uint32, base_shape, i)
           for i in range(len(base_shape))]
  eyes = [eq(i1, i2) for i1, i2 in zip(iotas[:-1], iotas[1:])]
  result = convert_element_type_p.bind(
      _reduce(operator.and_, eyes), new_dtype=dtype, weak_type=False,
      sharding=None)
  return broadcast_in_dim(result, shape, axes)

def _tri(dtype: DTypeLike, shape: Shape, offset: DimSize) -> Array:
  """Like numpy.tri, create a 2D array with ones below a diagonal."""
  offset = asarray(core.dimension_as_value(offset))
  if not dtypes.issubdtype(offset.dtype, np.integer):
    raise TypeError(f"offset must be an integer, got {offset!r}")
  shape_dtype = lax_utils.int_dtype_for_shape(shape, signed=True)
  if (
      np.iinfo(offset.dtype).min < np.iinfo(shape_dtype).min
      or np.iinfo(offset.dtype).max > np.iinfo(shape_dtype).max
  ):
    shape_dtype = np.dtype(np.int64)
  dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "tri")
  bool_tri = ge(add(broadcasted_iota(shape_dtype, shape, 0),
                    offset.astype(shape_dtype)),
                broadcasted_iota(shape_dtype, shape, 1))
  return convert_element_type_p.bind(bool_tri, new_dtype=dtype, weak_type=False,
                                     sharding=None)

def _stop_gradient(x):
  if dtypes.issubdtype(core.typeof(x).dtype, dtypes.extended):
    return x
  elif isinstance(x, ad.JVPTracer):
    return _stop_gradient(x.primal)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the offset to int before calling: jnp.tri(n, k=int(k)) or lax-level offset=np.int64(k)
2. Ensure traced offsets are produced by integer ops (use // not / when halving)
3. Type-check/validate external inputs that feed the offset parameter

Example fix

// before
k = n / 2
mask = jnp.tri(n, k=k)
// after
k = n // 2
mask = jnp.tri(n, k=k)

offset = int(offset) if not hasattr(offset, 'dtype') else offset
mask = jnp.tri(n, M, k=offset)

def is_int_like(v) -> bool:
    import numpy as np
    return np.issubdtype(getattr(v, 'dtype', type(v)), np.integer)

When it happens

Trigger: Calling jnp.tri(N, M, k=offset) or lax internally with a float offset (k=1.0), a Python float traced into asarray, or a weak-typed float scalar under jit.

Common situations: Computing k arithmetically (k = n/2 which yields float), passing user config values typed as float, or tracing jnp.tri under jit where offset comes from a float parameter.

Related errors

• full must be called with scalar fill_value, got fill_value.s
• {} does not accept dtype {}. Accepted dtypes are subtypes of
• {} does not accept dtype {} at position {}. Accepted dtypes
• the first argument to pow must have an inexact dtype (float
• Integers cannot be raised to negative powers, got integer_po

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