{"record":{"id":"260d87d37f42c2d4","repo":"jax-ml/jax","slug":"not-implemented-accuracy-260d87","errorCode":null,"errorMessage":"Not implemented: accuracy","messagePattern":"Not implemented: accuracy","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":2957,"sourceCode":"    return arith_dialect.mulf(x, x)\n  raise NotImplementedError(f\"Unsupported dtype {x_aval.dtype}\")\n\n\n@register_lowering_rule(lax.clz_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(lax.clz_p, mgpu.LoweringSemantics.Warpgroup)\ndef _clz_lowering_rule(ctx: LoweringRuleContext, x):\n  [x_aval] = ctx.avals_in\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n    return _ensure_fa(x, x_aval.dtype)._pointwise(math_dialect.ctlz, restrict_bitwidth=False)\n  x = _ensure_ir_value(x, x_aval.dtype)\n  return math_dialect.ctlz(x)\n\n\n@register_lowering_rule(lax.rsqrt_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(lax.rsqrt_p, mgpu.LoweringSemantics.Warpgroup)\ndef _rsqrt_lowering_rule(ctx: LoweringRuleContext, x, accuracy):\n  if accuracy is not None:\n    raise NotImplementedError(\"Not implemented: accuracy\")\n  [x_aval] = ctx.avals_in\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n    return _ensure_fa(x, x_aval.dtype).rsqrt(approx=ctx.module_ctx.approx_math)\n  fastmath = (\n      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None\n  )\n  return math_dialect.rsqrt(\n      _ensure_ir_value(x, x_aval.dtype), fastmath=fastmath\n  )\n\n\n@register_lowering_rule(lax.tanh_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(lax.tanh_p, mgpu.LoweringSemantics.Warpgroup)\ndef _tanh_lowering_rule(ctx: LoweringRuleContext, x, accuracy):\n  if accuracy is not None:\n    raise NotImplementedError(\"Not implemented: accuracy\")\n  [x_aval] = ctx.avals_in\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:","sourceCodeStart":2939,"sourceCodeEnd":2975,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L2939-L2975","documentation":"The Mosaic GPU lowering rule for lax.rsqrt accepts an accuracy hint argument but does not implement it; passing any non-None accuracy raises NotImplementedError('Not implemented: accuracy').","triggerScenarios":"Calling lax.rsqrt(x, accuracy=...) (or an API that forwards an accuracy hint, e.g. some precision-controlled math frontends) inside a Mosaic GPU kernel.","commonSituations":"Using accuracy-controlled math APIs or porting code that requests approximate vs precise rsqrt semantics; the kernel compiler path rejects the accuracy contract entirely.","solutions":["Call lax.rsqrt without the accuracy argument and rely on ctx.module_ctx.approx_math to control fastmath behavior","Set the approx_math flag on the compilation context instead of per-op accuracy","Fall back to 1 / jnp.sqrt(x) if specific accuracy behavior is needed"],"exampleFix":"// before\nr = lax.rsqrt(x, accuracy=1e-6)\n// after\nr = lax.rsqrt(x)  # fastmath controlled by module approx_math","handlingStrategy":"validation","validationCode":"# strip accuracy hints before kernel launch\ndef rsqrt_safe(x, accuracy=None):\n  if accuracy is not None:\n    raise ValueError('accuracy unsupported in Mosaic; use approx_math flag')\n  return lax.rsqrt(x)","typeGuard":null,"tryCatchPattern":"try:\n  out = kernel(x)\nexcept NotImplementedError as e:\n  if 'accuracy' in str(e): call lax.rsqrt(x) without accuracy or use 1/jnp.sqrt","preventionTips":["Never pass accuracy to lax math inside Pallas kernels","Control precision globally via module approx_math instead"],"tags":["jax","pallas","mosaic-gpu","rsqrt","accuracy","not-implemented"],"backgroundTag":"unsupported-parameter-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}