{"record":{"id":"481e61a921b027c0","repo":"jax-ml/jax","slug":"unsupported-dtypes-x-aval-dtype-and-y-aval-dty","errorCode":null,"errorMessage":"unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}","messagePattern":"unsupported dtypes: (.+?) and (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1427,"sourceCode":"    lax.clamp_p: lambda min, a, max: jnp.minimum(jnp.maximum(min, a), max),\n    lax.logistic_p: lambda a, accuracy: 1 / (1 + jnp.exp(-a)),\n    lax.is_finite_p: lambda x: jnp.logical_and(~jnp.isnan(x), ~jnp.isinf(x)),\n}\n\nfor prim, fn in _JAX_FN_MAPPING.items():\n  triton_lowering_rules[prim] = lower_fun(fn, multiple_results=False)\n\n\n@register_lowering(lax.min_p)\ndef _min_lowering_rule(ctx: LoweringRuleContext, x, y):\n  # TODO(slebedev): Consider allowing customizing nan behavior.\n  x_aval, y_aval = ctx.avals_in\n  x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)\n  if jnp.issubdtype(x_aval.dtype, jnp.floating):\n    # TODO(slebedev): Triton promotes bfloat16 to float32 and back here.\n    return arith_dialect.minnumf(x, y)\n  if not jnp.issubdtype(x_aval.dtype, jnp.integer):\n    raise NotImplementedError(\n        f\"unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}\"\n    )\n  if jnp.issubdtype(x_aval.dtype, jnp.signedinteger):\n    return arith_dialect.minsi(x, y)\n  else:\n    return arith_dialect.minui(x, y)\n\n\n@register_lowering(lax.max_p)\ndef _max_lowering_rule(ctx: LoweringRuleContext, x, y):\n  # TODO(slebedev): Consider allowing customizing nan behavior.\n  x_aval, y_aval = ctx.avals_in\n  x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)\n  if jnp.issubdtype(x_aval.dtype, jnp.floating):\n    # TODO(slebedev): Triton promotes bfloat16 to float32 and back here.\n    return arith_dialect.maxnumf(x, y)\n  if not jnp.issubdtype(x_aval.dtype, jnp.integer):\n    raise NotImplementedError(","sourceCodeStart":1409,"sourceCodeEnd":1445,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1409-L1445","documentation":"The Triton lowering rule for jnp.minimum only implements floating-point and integer element types. If the inputs' dtype is neither (e.g. complex, bool, or a custom dtype), lowering fails with this NotImplementedError. It reflects Triton's arith dialect ops (minnumf/minsi/minui) which only cover those categories.","triggerScenarios":"Calling jnp.minimum(x, y) inside a Pallas kernel lowered to Triton where x_aval.dtype is complex64/complex128, bool_, or a non-numeric dtype; also float8 dtypes that are not yet handled.","commonSituations":"Writing Pallas kernels that operate on complex numbers (e.g. FFT-like code) or boolean masks using minimum instead of logical_and; migrating code from XLA/TPU where minimum on these dtypes works; enabling float8 experiments on newer JAX versions.","solutions":["Convert inputs to a supported dtype before the call: use jnp.where or replace jnp.minimum on bools with jnp.logical_and / jnp.bitwise_and","For complex inputs, split into real/imag parts, apply minimum to each, and recombine, or compute the comparison manually with real()-based logic","Check jax/issues for float8/complex min support in the Triton backend and upgrade JAX if support landed"],"exampleFix":"// before\nm = jnp.minimum(mask_a, mask_b)  # bool inputs\n// after\nm = jnp.logical_and(mask_a, mask_b)","handlingStrategy":"type-guard","validationCode":"SUPPORTED = lambda d: jnp.issubdtype(d, jnp.floating) or jnp.issubdtype(d, jnp.integer)\nassert SUPPORTED(x.dtype) and SUPPORTED(y.dtype), 'minimum on Triton needs float/int'","typeGuard":"def triton_min_safe(x, y):\n    if x.dtype == jnp.bool_ or y.dtype == jnp.bool_:\n        return jnp.logical_and(x, y)\n    if jnp.issubdtype(x.dtype, jnp.complexfloating):\n        raise NotImplementedError('complex minimum unsupported on Triton')\n    return jnp.minimum(x, y)","tryCatchPattern":"try:\n    kernel = pallas.triton_compile(...)  # or jitted call\nexcept NotImplementedError as e:\n    if 'unsupported dtypes' in str(e):\n        # fall back to an XLA-jitted equivalent\n        out = jax.jit(jnp.minimum)(x, y)\n    else:\n        raise","preventionTips":["Keep kernel dtypes to standard float/int types; use logical_and/or for booleans","Add dtype asserts at kernel entry to fail fast with a clear message"],"tags":["jax","pallas","triton","jnp-minimum","dtype","notimplementederror"],"backgroundTag":"unsupported-dtype-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}