{"record":{"id":"12bdc21a4a95abf0","repo":"jax-ml/jax","slug":"ctx-prim-does-not-support-x-aval-dtype-and-y","errorCode":null,"errorMessage":"{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype}","messagePattern":"(.+?) does not support (.+?) and (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":2789,"sourceCode":"    ctx: LoweringRuleContext, x, y, *, ui_impl, si_impl, f_impl=None, **kwargs,\n):\n  if kwargs.get('out_dtype') is not None:\n    raise NotImplementedError(\"out_dtype argument in binary_op_lowering_rule_wg\")\n  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n    if any(aval_in.shape for aval_in in ctx.avals_in):\n      raise NotImplementedError(\n          \"Non-scalar arithmetic is not supported in warp-level lowering.\")\n  x_aval, y_aval = ctx.avals_in\n  [out_aval] = ctx.avals_out\n  x, y = _bcast_wg(x, y, *ctx.avals_in, *ctx.avals_out)\n  if jnp.issubdtype(out_aval, jnp.signedinteger):\n    return si_impl(x, y)\n  elif jnp.issubdtype(out_aval, jnp.integer):\n    return ui_impl(x, y)\n  elif f_impl is not None and jnp.issubdtype(out_aval, jnp.floating):\n    return f_impl(x, y)\n  else:\n    raise NotImplementedError(\n        f\"{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype}\"\n    )\n\n\nfor op, si_impl, ui_impl, f_impl in [\n    (lax.add_p, arith_dialect.addi, arith_dialect.addi, arith_dialect.addf),\n    (lax.sub_p, arith_dialect.subi, arith_dialect.subi, arith_dialect.subf),\n    (lax.mul_p, arith_dialect.muli, arith_dialect.muli, arith_dialect.mulf),\n    (\n        lax.div_p,\n        arith_dialect.divsi,\n        arith_dialect.divui,\n        arith_dialect.divf,\n    ),\n    (lax.rem_p, arith_dialect.remsi, arith_dialect.remui, arith_dialect.remf),\n    (\n        lax.max_p,\n        arith_dialect.maxsi,","sourceCodeStart":2771,"sourceCodeEnd":2807,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L2771-L2807","documentation":"The binary op lowering dispatches on the output aval's dtype class: signed integer -> addi/subi/muli, unsigned integer -> same integer ops, floating -> addf/subf/mulf (only when f_impl is provided). If the output dtype matches none of these (e.g. complex, bool, bfloat16 edge cases, or an op lacking f_impl), it raises NotImplementedError naming the primitive and both input dtypes.","triggerScenarios":"Applying a lax binary primitive inside a Mosaic GPU kernel to operands whose promoted dtype is complex (or another unsupported class), or a float dtype for an op registered without an f_impl, e.g. integer-only ops on floats.","commonSituations":"Using complex128/complex64 math or unsupported float types (f8 etc.) in Pallas kernels; applying ops like shift_left or rem to floating inputs where no f_impl exists.","solutions":["Cast inputs to a supported dtype (float32/float64 or int32/int64) before the op","Avoid complex dtypes inside Mosaic GPU kernels; implement the op on real/imag parts manually","Check the registration table in lowering.py to confirm the op supports your dtype class","Use a non-Pallas implementation for unsupported dtype/op combinations"],"exampleFix":"// before\nz = a_complex * b_complex  # complex dtype unsupported\n// after\nz_re = a.real * b.real - a.imag * b.imag\nz_im = a.real * b.imag + a.imag * b.real","handlingStrategy":"type-guard","validationCode":"import jax.numpy as jnp\nSUPPORTED = (jnp.integer, jnp.floating)\ndef ok_dtype(*xs):\n  return all(jnp.issubdtype(x.dtype, SUPPORTED) for x in xs)","typeGuard":"def is_supported_binary_dtype(x, y):\n  d = jnp.promote_types(x.dtype, y.dtype)\n  return jnp.issubdtype(d, jnp.integer) or jnp.issubdtype(d, jnp.floating)","tryCatchPattern":"try:\n  out = kernel(x, y)\nexcept NotImplementedError as e:\n  if 'does not support' in str(e): cast inputs to float32/int32 and retry","preventionTips":["Standardize kernel inputs to float32/int32","Avoid complex dtypes in Mosaic kernels"],"tags":["jax","pallas","mosaic-gpu","dtype","binary-op","not-implemented"],"backgroundTag":"unsupported-dtype-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}