{"record":{"id":"62e0c63318b95379","repo":"jax-ml/jax","slug":"logical-reduction-requires-operand-dtype-bool-or-i","errorCode":null,"errorMessage":"logical reduction requires operand dtype bool or int, got {operand.dtype}.","messagePattern":"logical reduction requires operand dtype bool or int, got (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8748,"sourceCode":"    mlir.lower_fun(\n        partial(_compute_argminmax, lt, _get_min_identity),\n        multiple_results=False,\n    ),\n    inline=False,\n)\n\nmlir.register_lowering(\n    argmax_p,\n    mlir.lower_fun(\n        partial(_compute_argminmax, gt, _get_max_identity),\n        multiple_results=False,\n    ),\n    inline=False,\n)\n\ndef _reduce_logical_shape_rule(operand, *, axes):\n  if operand.dtype != np.bool_ and not np.issubdtype(operand.dtype, np.integer):\n    raise TypeError(f\"logical reduction requires operand dtype bool or int, got {operand.dtype}.\")\n  return tuple(np.delete(operand.shape, axes))\n\ndef _reduce_logical_sharding_rule(operand, *, axes):\n  return operand.sharding.update(spec=tuple_delete(operand.sharding.spec, axes))\n\ndef _reduce_or_lin(_is_vjp, nzs, x, *, axes):\n  nz, = nzs\n  y = reduce_or_p.bind(x, axes=axes)\n  aval = typeof(y).to_tangent_aval()\n  return y, False, (), None, lambda _, __, t: ad_util.Zero(aval)\n\nreduce_or_p = standard_primitive(\n    _reduce_logical_shape_rule, input_dtype, 'reduce_or',\n    weak_type_rule=_strip_weak_type, sharding_rule=_reduce_logical_sharding_rule,\n    vma_rule=partial(core.standard_vma_rule, 'reduce_or'))\nbatching.defreducer(reduce_or_p)\nad.primitive_linearizations[reduce_or_p] = _reduce_or_lin\n","sourceCodeStart":8730,"sourceCodeEnd":8766,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8730-L8766","documentation":"Logical reductions (reduce_and, reduce_or, reduce_xor) require the operand to have bool or integer dtype; float or complex operands are rejected because bitwise logical ops are undefined for them.","triggerScenarios":"lax.reduce_and(x_float, axes), jnp.logical_and-style reductions applied via lax on a float array, or calling the internal _reduce_or_lin transpose path with non-integer input.","commonSituations":"Applying all()/any()-style reductions implemented with bitwise ops to boolean masks that were accidentally cast to float (e.g., after arithmetic like mask * 1.0); data-type drift in a preprocessing pipeline.","solutions":["Cast the operand to bool before the reduction: lax.reduce_or(x != 0, axes) or x.astype(jnp.bool_).","Fix upstream so the mask stays boolean instead of being promoted to float.","Use jnp.all / jnp.any for logical semantics on any dtype-comparable input."],"exampleFix":"# before\nout = lax.reduce_or(mask * 1.0, (0,))\n# after\nout = lax.reduce_or((mask * 1.0).astype(jnp.bool_), (0,))","handlingStrategy":"type-guard","validationCode":"if x.dtype != jnp.bool_:\n    x = x != 0\nout = lax.reduce_or(x, (0,))","typeGuard":"def bool_or_int_dtype(x):\n    import numpy as np\n    return x.dtype == np.bool_ or np.issubdtype(x.dtype, np.integer)","tryCatchPattern":null,"preventionTips":["Keep masks boolean; avoid arithmetic that promotes masks to float.","Use jnp.any/jnp.all for logical semantics."],"tags":["jax","lax","logical-reduction","dtype"],"backgroundTag":"invalid-dtype-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}