{"record":{"id":"b20896df18d246b8","repo":"jax-ml/jax","slug":"offset-must-be-an-integer-got-offset-r","errorCode":null,"errorMessage":"offset must be an integer, got {offset!r}","messagePattern":"offset must be an integer, got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":3709,"sourceCode":"\ndef _delta(dtype: DTypeLike, shape: Shape, axes: Sequence[int]) -> Array:\n  \"\"\"This utility function exists for creating Kronecker delta arrays.\"\"\"\n  axes = map(int, axes)\n  dtype = dtypes.check_and_canonicalize_user_dtype(dtype, \"delta\")\n  base_shape = tuple(np.take(shape, axes))\n  iotas = [broadcasted_iota(np.uint32, base_shape, i)\n           for i in range(len(base_shape))]\n  eyes = [eq(i1, i2) for i1, i2 in zip(iotas[:-1], iotas[1:])]\n  result = convert_element_type_p.bind(\n      _reduce(operator.and_, eyes), new_dtype=dtype, weak_type=False,\n      sharding=None)\n  return broadcast_in_dim(result, shape, axes)\n\ndef _tri(dtype: DTypeLike, shape: Shape, offset: DimSize) -> Array:\n  \"\"\"Like numpy.tri, create a 2D array with ones below a diagonal.\"\"\"\n  offset = asarray(core.dimension_as_value(offset))\n  if not dtypes.issubdtype(offset.dtype, np.integer):\n    raise TypeError(f\"offset must be an integer, got {offset!r}\")\n  shape_dtype = lax_utils.int_dtype_for_shape(shape, signed=True)\n  if (\n      np.iinfo(offset.dtype).min < np.iinfo(shape_dtype).min\n      or np.iinfo(offset.dtype).max > np.iinfo(shape_dtype).max\n  ):\n    shape_dtype = np.dtype(np.int64)\n  dtype = dtypes.check_and_canonicalize_user_dtype(dtype, \"tri\")\n  bool_tri = ge(add(broadcasted_iota(shape_dtype, shape, 0),\n                    offset.astype(shape_dtype)),\n                broadcasted_iota(shape_dtype, shape, 1))\n  return convert_element_type_p.bind(bool_tri, new_dtype=dtype, weak_type=False,\n                                     sharding=None)\n\ndef _stop_gradient(x):\n  if dtypes.issubdtype(core.typeof(x).dtype, dtypes.extended):\n    return x\n  elif isinstance(x, ad.JVPTracer):\n    return _stop_gradient(x.primal)","sourceCodeStart":3691,"sourceCodeEnd":3727,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L3691-L3727","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Cast the offset to int before calling: jnp.tri(n, k=int(k)) or lax-level offset=np.int64(k)","Ensure traced offsets are produced by integer ops (use // not / when halving)","Type-check/validate external inputs that feed the offset parameter"],"exampleFix":"// before\nk = n / 2\nmask = jnp.tri(n, k=k)\n// after\nk = n // 2\nmask = jnp.tri(n, k=k)","handlingStrategy":"type-guard","validationCode":"offset = int(offset) if not hasattr(offset, 'dtype') else offset\nmask = jnp.tri(n, M, k=offset)","typeGuard":"def is_int_like(v) -> bool:\n    import numpy as np\n    return np.issubdtype(getattr(v, 'dtype', type(v)), np.integer)","tryCatchPattern":null,"preventionTips":["Keep diagonal offsets as ints; use // not / when computing them","Cast user-supplied k with int() before jnp.tri"],"tags":["jax","lax","tri","integer-required","type-error"],"backgroundTag":"wrong-dtype-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}