jax-ml/jax · error · NotImplementedError

erf_inv_lowering_helper not implemented for

Error message

erf_inv_lowering_helper not implemented for {x.dtype}

What it means

pallas.utils.erf_inv_lowering_helper implements inverse error function only for float32 and float64 inputs. Any other dtype (bfloat16, float16, integers) raises NotImplementedError because no polynomial approximation is wired up for it.

Solutions

1. Cast to float32 before calling: x.astype(jnp.float32), cast back after
2. Compute erfinv outside the pallas kernel on the host/XLA side

Example fix

# before
y = jax.scipy.special.erfinv(x_bf16)
# after
y = jax.scipy.special.erfinv(x_bf16.astype(jnp.float32)).astype(x_bf16.dtype)

assert x.dtype in (jnp.float32, jnp.float64)
x32 = x.astype(jnp.float32)

When it happens

Trigger: Calling jax.scipy.special.erfinv (or erf_inv) inside a pallas triton kernel on a bfloat16/float16/int array.

Common situations: Kernels written with autocast/bfloat16 defaults (common on TPUs/GPUs) that call erfinv; sampling code (e.g. truncated normal via inverse CDF) in low precision.

Related errors

• a and b must have the same element type, but got
• cannot cast from
• cannot cast to
• cmp and val must have identical dtypes and shapes
• `nextafter` only supports float32 and float64, but got

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

Appendix: source

Thrown at jax/_src/pallas/utils.py:349

  w = jnp.where(w_lt_625, w - 3.125, select2_result)

  p = get_coefficient(0)
  for i in range(1, 17):
    p = get_coefficient(i) + p * w
  for i in range(17, 19):
    p = jnp.where(w_lt_16, get_coefficient(i) + p * w, p)
  for i in range(19, 23):
    p = jnp.where(w_lt_625, get_coefficient(i) + p * w, p)

  return jnp.where(jnp.abs(x) == 1.0, np.inf * x, p * x)


def erf_inv_lowering_helper(x):
  if x.dtype == jnp.float32:
    return _erf_inv_32_lowering_helper(x)
  if x.dtype == jnp.float64:
    return _erf_inv_64_lowering_helper(x)
  raise NotImplementedError(f"erf_inv_lowering_helper not implemented for {x.dtype}")


def sign_lowering_helper(x):
  if jnp.issubdtype(x.dtype, jnp.unsignedinteger):
    return (x != 0).astype(x.dtype)

  if jnp.issubdtype(x.dtype, jnp.integer):
    return (x > 0).astype(x.dtype) - (x < 0).astype(x.dtype)

  if jnp.issubdtype(x.dtype, jnp.floating):
    out = (x > 0.).astype(x.dtype) - (x < 0.).astype(x.dtype)
    return jnp.where(jnp.isnan(x), jnp.nan, out)

  raise NotImplementedError(f"sign_lowering_helper not implemented for {x.dtype}")


# based on https://github.com/openxla/xla/blob/a7a09d56c3599123f8148bbf3e44c9ebc04624b9/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo.cc#L1339-L1422
def nextafter_lowering_helper(x, y):

View on GitHub (pinned to 1e1c6a8fc0)