{"record":{"id":"46dc72779eff455e","repo":"jax-ml/jax","slug":"unsupported-dtype-x-aval-dtype","errorCode":null,"errorMessage":"Unsupported dtype {x_aval.dtype}","messagePattern":"Unsupported dtype (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":2911,"sourceCode":"\n@register_lowering_rule(lax.integer_pow_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(lax.integer_pow_p, mgpu.LoweringSemantics.Warpgroup)\ndef _integer_pow_lowering_rule(ctx: LoweringRuleContext, x, y):\n  [x_aval] = ctx.avals_in\n  if y == -1:\n    return _lower_fun(lambda x: 1 / x)(ctx, x)\n  if y <= 1:\n    raise NotImplementedError\n\n  mul_op: Callable[[Any, Any], Any]\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n    mul_op = operator.mul\n  elif jnp.issubdtype(x_aval.dtype, jnp.integer):\n    mul_op = arith_dialect.muli\n  elif jnp.issubdtype(x_aval.dtype, jnp.floating):\n    mul_op = arith_dialect.mulf\n  else:\n    raise NotImplementedError(f\"Unsupported dtype {x_aval.dtype}\")\n\n  # Y is an integer. Here we start with res = x so the range is y-1\n  res = x\n  # Repeated doubling algorithm.\n  for i in reversed(range(y.bit_length() - 1)):\n    res = mul_op(res, res)  # pyrefly: ignore[no-matching-overload]\n    if (y >> i) & 1:\n      res = mul_op(res, x)\n  return res\n\n\n@register_lowering_rule(lax.clamp_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(lax.clamp_p, mgpu.LoweringSemantics.Warpgroup)\ndef _clamp_lowering_rule(ctx: LoweringRuleContext, l, x, u):\n  return _lower_fun(lambda l, x, u: lax.min(lax.max(x, l), u))(ctx, l, x, u)\n\n\n@register_lowering_rule(lax.square_p, mgpu.LoweringSemantics.Lane)","sourceCodeStart":2893,"sourceCodeEnd":2929,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L2893-L2929","documentation":"In integer_pow's repeated-squaring loop, the multiplication op is chosen from the input dtype: integer -> arith.muli, floating -> arith.mulf. Any other dtype class (complex, bool) has no multiplier and raises 'Unsupported dtype {x_aval.dtype}'.","triggerScenarios":"Calling lax.integer_pow with y >= 2 on complex (or other non-int/float) operands inside a Mosaic GPU kernel.","commonSituations":"Computing integer powers of complex numbers (e.g. z**2, z**3) in a Pallas kernel; complex arithmetic lacks a direct MLIR mul mapping here.","solutions":["Expand the power manually into real/imaginary multiply-adds","Cast/keep inputs float or integer for power operations","Compute complex powers outside the kernel with standard JAX"],"exampleFix":"// before\nz2 = z ** 2  # complex -> Unsupported dtype\n// after\nz2_re = z.real*z.real - z.imag*z.imag\nz2_im = 2*z.real*z.imag","handlingStrategy":"type-guard","validationCode":"import jax.numpy as jnp\n# only int/float bases support repeated-squaring powers\ndef pow_dtype_ok(x):\n  return jnp.issubdtype(x.dtype, jnp.integer) or jnp.issubdtype(x.dtype, jnp.floating)","typeGuard":"def supports_integer_pow(x):\n  return (jnp.issubdtype(x.dtype, jnp.integer)\n          or jnp.issubdtype(x.dtype, jnp.floating))","tryCatchPattern":"try:\n  out = kernel(x)\nexcept NotImplementedError as e:\n  if 'Unsupported dtype' in str(e): expand complex power manually","preventionTips":["Use real/imag decomposition for complex powers","Keep power operands float32 where possible"],"tags":["jax","pallas","mosaic-gpu","integer-pow","complex-dtype","not-implemented"],"backgroundTag":"unsupported-dtype-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}