{"record":{"id":"180aec22cec91848","repo":"jax-ml/jax","slug":"called-name-with-a-float0-at-position-i-float","errorCode":null,"errorMessage":"Called {name} with a float0 at position {i}. float0s do not support any operations by design, because they are not compatible with non-trivial vector spaces. No implicit dtype conversion is done. You can use np.zeros_like(arr, dtype=np.float) to cast a float0 array to a regular zeros array. \nIf you didn't expect to get a float0 you might have accidentally taken a gradient with respect to an integer argument.","messagePattern":"Called (.+?) with a float0 at position (.+?)\\. float0s do not support any operations by design, because they are not compatible with non-trivial vector spaces\\. No implicit dtype conversion is done\\. You can use np\\.zeros_like\\(arr, dtype=np\\.float\\) to cast a float0 array to a regular zeros array\\. \nIf you didn't expect to get a float0 you might have accidentally taken a gradient with respect to an integer argument\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":4305,"sourceCode":"                            ur_rule=partial(unop_ur_rule, name))\n  batching.defvectorized(prim)\n  return prim\n\nstandard_unop = partial(unop, _identity)\n\n_attrgetter = lambda name: lambda x, **kwargs: getattr(x, name)\n\n\ndef naryop_dtype_rule(result_dtype, accepted_dtypes, name, *avals,\n                      require_same=True, allow_extended_dtype=False, **kwargs):\n  assert len(avals) == len(accepted_dtypes), (avals, accepted_dtypes)\n  for i, aval in enumerate(avals):\n    if allow_extended_dtype and isinstance(aval.dtype, dtypes.ExtendedDType):\n      continue\n    types = accepted_dtypes[i]\n    if not any(dtypes.issubdtype(aval.dtype, t) for t in types):\n      if aval.dtype == dtypes.float0:\n        raise TypeError(\n            f\"Called {name} with a float0 at position {i}. \"\n            \"float0s do not support any operations by design, because they \"\n            \"are not compatible with non-trivial vector spaces. No implicit dtype \"\n            \"conversion is done. You can use np.zeros_like(arr, dtype=np.float) \"\n            \"to cast a float0 array to a regular zeros array. \\n\"\n            \"If you didn't expect to get a float0 you might have accidentally \"\n            \"taken a gradient with respect to an integer argument.\")\n      else:\n        msg = ('{} does not accept dtype {} at position {}. '\n               'Accepted dtypes at position {} are subtypes of {}.')\n        typename = dtype_to_string(aval.dtype)\n        typenames = ', '.join(t.__name__ for t in types)\n        raise TypeError(msg.format(name, typename, i, i, typenames))\n  if require_same and kwargs.get('out_dtype') is None:\n    check_same_dtypes(name, *avals)\n  return result_dtype(*avals, **kwargs)\n\n","sourceCodeStart":4287,"sourceCodeEnd":4323,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L4287-L4323","documentation":"N-ary lax ops (binops like add, mul) validate each operand position against per-position accepted dtypes. If the offending operand has dtype float0 — the tangent dtype of non-differentiable values — this specialized TypeError is raised instead of the generic dtype message, since float0 typically signals an autodiff misuse.","triggerScenarios":"Using jvp/vjp/grad where a float0 tangent (from an integer or boolean primal) flows into a binary lax op at position i; chaining integer inputs through differentiable code under transformation.","commonSituations":"grad/jvp of functions mixing integer indices with float math; custom_vjp rules that pass raw tangents (possibly float0) into lax ops; masking pipelines where boolean masks become float0 tangents.","solutions":["stop_gradient integer/boolean inputs so their tangents don't propagate: lax.stop_gradient(idx)","Cast integer primals to float before differentiating","Register custom VJPs that return zero tangents for non-differentiable inputs"],"exampleFix":"// before\ndef f(x, idx):\n    return (x * w[idx]).sum()\njax.grad(f, argnums=1)(x, idx)  # float0 tangent into mul\n// after\ndef f(x, idx):\n    return (x * w[lax.stop_gradient(idx)]).sum()","handlingStrategy":"type-guard","validationCode":"import jax.numpy as jnp, jax.dtypes as dt\nif any(a.dtype == dt.float0 for a in (x, y)):\n    x = jnp.zeros_like(x, dtype=jnp.float32) if x.dtype == dt.float0 else x\n    y = jnp.zeros_like(y, dtype=jnp.float32) if y.dtype == dt.float0 else y","typeGuard":"def any_float0(*args) -> bool:\n    dt = __import__('jax').dtypes\n    return any(a.dtype == dt.float0 for a in args)","tryCatchPattern":null,"preventionTips":["stop_gradient integer/boolean operands in differentiable functions","Keep integer indices out of argnums of grad/jacfwd","Cast to float any quantity you intend to differentiate through"],"tags":["jax","float0","autodiff","binop","type-error"],"backgroundTag":"gradient-of-integer-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}