jax-ml/jax · error · NotImplementedError

Not implemented: accuracy

Error message

Not implemented: accuracy

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3842

    raise ValueError(f"Unsupported dtype for sign: {x.dtype}")

  return lower_fun(_lower_fun)(ctx, x)


@register_lowering_rule(lax.nextafter_p)
def _nextafter_lowering_rule(ctx: LoweringRuleContext, x, y):
  return lower_fun(
      pallas_utils.nextafter_lowering_helper,
  )(ctx, x, y)


@register_lowering_rule(
    lax.rsqrt_p,
    kernel_types=(tpu_core.CoreType.TC, tpu_core.CoreType.SC_VECTOR_SUBCORE),
)
def _rsqrt_lowering_rule(ctx: LoweringRuleContext, x, accuracy=None):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  return mlir_math.rsqrt(x)


@register_lowering_rule(
    lax.sqrt_p,
    kernel_types=(tpu_core.CoreType.TC, tpu_core.CoreType.SC_VECTOR_SUBCORE),
)
def _sqrt_lowering_rule(ctx: LoweringRuleContext, x, accuracy=None):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  return mlir_math.sqrt(x)


@register_lowering_rule(lax.square_p)
def _square_lowering_rule(ctx: LoweringRuleContext, x):
  if jnp.issubdtype(ctx.avals_in[0].dtype, jnp.integer):
    return arith.muli(x, x)
  return arith.mulf(x, x)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call lax.rsqrt without the accuracy argument
2. Use 1 / jnp.sqrt(x) if you need to control precision manually
3. Strip accuracy hints when targeting TPU Pallas kernels

Example fix

// before
r = lax.rsqrt(x, accuracy=0.01)
// after
r = lax.rsqrt(x)

def rsqrt_call_safe(accuracy):
    return accuracy is None

When it happens

Trigger: 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).

Common situations: Code that sets accuracy hints for XLA CPU/GPU fast-math reused on TPU; library code (e.g. attention implementations) parameterizing rsqrt accuracy.

Related errors

• Not implemented: accuracy
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/0bafe2db41af678e. Report an issue: GitHub.