{"record":{"id":"0bafe2db41af678e","repo":"jax-ml/jax","slug":"not-implemented-accuracy","errorCode":null,"errorMessage":"Not implemented: accuracy","messagePattern":"Not implemented: accuracy","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3842,"sourceCode":"    raise ValueError(f\"Unsupported dtype for sign: {x.dtype}\")\n\n  return lower_fun(_lower_fun)(ctx, x)\n\n\n@register_lowering_rule(lax.nextafter_p)\ndef _nextafter_lowering_rule(ctx: LoweringRuleContext, x, y):\n  return lower_fun(\n      pallas_utils.nextafter_lowering_helper,\n  )(ctx, x, y)\n\n\n@register_lowering_rule(\n    lax.rsqrt_p,\n    kernel_types=(tpu_core.CoreType.TC, tpu_core.CoreType.SC_VECTOR_SUBCORE),\n)\ndef _rsqrt_lowering_rule(ctx: LoweringRuleContext, x, accuracy=None):\n  if accuracy is not None:\n    raise NotImplementedError(\"Not implemented: accuracy\")\n  return mlir_math.rsqrt(x)\n\n\n@register_lowering_rule(\n    lax.sqrt_p,\n    kernel_types=(tpu_core.CoreType.TC, tpu_core.CoreType.SC_VECTOR_SUBCORE),\n)\ndef _sqrt_lowering_rule(ctx: LoweringRuleContext, x, accuracy=None):\n  if accuracy is not None:\n    raise NotImplementedError(\"Not implemented: accuracy\")\n  return mlir_math.sqrt(x)\n\n\n@register_lowering_rule(lax.square_p)\ndef _square_lowering_rule(ctx: LoweringRuleContext, x):\n  if jnp.issubdtype(ctx.avals_in[0].dtype, jnp.integer):\n    return arith.muli(x, x)\n  return arith.mulf(x, x)","sourceCodeStart":3824,"sourceCodeEnd":3860,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3824-L3860","documentation":"The Mosaic lowering for lax.rsqrt (on TensorCore and SC_VECTOR_SUBCORE kernels) does not implement the optional accuracy parameter. Passing any non-None accuracy (used for fast approximate math on other backends) raises NotImplementedError.","triggerScenarios":"Calling lax.rsqrt(x, accuracy=...) with an accuracy specification inside a Pallas TPU kernel, often indirectly via precision configs or libraries that thread accuracy hints (e.g. some training frameworks' fast-math paths).","commonSituations":"Code that sets accuracy hints for XLA CPU/GPU fast-math reused on TPU; library code (e.g. attention implementations) parameterizing rsqrt accuracy.","solutions":["Call lax.rsqrt without the accuracy argument","Use 1 / jnp.sqrt(x) if you need to control precision manually","Strip accuracy hints when targeting TPU Pallas kernels"],"exampleFix":"// before\nr = lax.rsqrt(x, accuracy=0.01)\n// after\nr = lax.rsqrt(x)","handlingStrategy":"validation","validationCode":"def rsqrt_call_safe(accuracy):\n    return accuracy is None","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never pass accuracy to lax.rsqrt in TPU Pallas kernels","Strip precision/accuracy hints from shared code paths targeting TPU"],"tags":["jax","pallas","tpu","rsqrt","accuracy"],"backgroundTag":"unsupported-parameter-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}