{"record":{"id":"9782d3585f94347d","repo":"jax-ml/jax","slug":"does-not-accept-dtype-at-position-accept","errorCode":null,"errorMessage":"{} does not accept dtype {} at position {}. Accepted dtypes at position {} are subtypes of {}.","messagePattern":"(.+?) does not accept dtype (.+?) at position (.+?)\\. Accepted dtypes at position (.+?) are subtypes of (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":4318,"sourceCode":"    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\ndef broadcasting_shape_rule(name, *avals, **kwargs):\n  if not isinstance(name, str):\n    raise RuntimeError(\n      \"First argument of broadcasting_shape_rule should be a name.\"\n      f\" Got {name}\")\n  shapes = [aval.shape for aval in avals if aval.shape]\n  if not shapes:\n    return ()\n  return _try_broadcast_shapes(*shapes, name=name)\n\n\ndef broadcasting_sharding_rule(name, *avals, **kwargs):\n  if not isinstance(name, str):","sourceCodeStart":4300,"sourceCodeEnd":4336,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L4300-L4336","documentation":"N-ary lax ops validate each argument position's dtype against a per-position accepted_dtypes list. When operand i's dtype is not a subtype of any accepted category (and isn't float0), this TypeError reports the position and the allowed types.","triggerScenarios":"Passing e.g. a complex operand where only real floats are accepted at that position, or a bool where integer/float required, to lax binops like pow's operands, comparison internals, or division variants.","commonSituations":"Complex-valued state accidentally fed into real-only ops; bool arrays used as numbers (JAX does not implicitly promote bool to int in lax); mixed dtypes across versions where accepted lists changed.","solutions":["Cast the offending argument to an accepted dtype: x.astype(jnp.float32)","Check the op signature/docstring for which dtypes each position accepts","Ensure arrays aren't bool when arithmetic dtypes are expected (use .astype(int) first)"],"exampleFix":"// before\nz = lax.integer_pow(bmask)  # bool at position 0\n// after\nz = lax.integer_pow(bmask.astype(jnp.int32))","handlingStrategy":"type-guard","validationCode":"x = x.astype(jnp.float32) if not jnp.issubdtype(x.dtype, jnp.inexact) else x\ny = y.astype(jnp.float32) if not jnp.issubdtype(y.dtype, jnp.inexact) else y","typeGuard":"def operands_acceptable(args, accepted) -> bool:\n    import numpy as np\n    return all(any(np.issubdtype(a.dtype, t) for t in ts)\n               for a, ts in zip(args, accepted))","tryCatchPattern":null,"preventionTips":["Centralize dtype normalization at data ingestion","Avoid bool arrays in arithmetic; cast to int32/float32 explicitly"],"tags":["jax","lax","binop","dtype-validation","type-error"],"backgroundTag":"unsupported-dtype-for-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}