{"record":{"id":"42a896334a57f682","repo":"jax-ml/jax","slug":"erf-inv-lowering-helper-not-implemented-for-x-dty","errorCode":null,"errorMessage":"erf_inv_lowering_helper not implemented for {x.dtype}","messagePattern":"erf_inv_lowering_helper not implemented for (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/utils.py","lineNumber":349,"sourceCode":"  w = jnp.where(w_lt_625, w - 3.125, select2_result)\n\n  p = get_coefficient(0)\n  for i in range(1, 17):\n    p = get_coefficient(i) + p * w\n  for i in range(17, 19):\n    p = jnp.where(w_lt_16, get_coefficient(i) + p * w, p)\n  for i in range(19, 23):\n    p = jnp.where(w_lt_625, get_coefficient(i) + p * w, p)\n\n  return jnp.where(jnp.abs(x) == 1.0, np.inf * x, p * x)\n\n\ndef erf_inv_lowering_helper(x):\n  if x.dtype == jnp.float32:\n    return _erf_inv_32_lowering_helper(x)\n  if x.dtype == jnp.float64:\n    return _erf_inv_64_lowering_helper(x)\n  raise NotImplementedError(f\"erf_inv_lowering_helper not implemented for {x.dtype}\")\n\n\ndef sign_lowering_helper(x):\n  if jnp.issubdtype(x.dtype, jnp.unsignedinteger):\n    return (x != 0).astype(x.dtype)\n\n  if jnp.issubdtype(x.dtype, jnp.integer):\n    return (x > 0).astype(x.dtype) - (x < 0).astype(x.dtype)\n\n  if jnp.issubdtype(x.dtype, jnp.floating):\n    out = (x > 0.).astype(x.dtype) - (x < 0.).astype(x.dtype)\n    return jnp.where(jnp.isnan(x), jnp.nan, out)\n\n  raise NotImplementedError(f\"sign_lowering_helper not implemented for {x.dtype}\")\n\n\n# 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\ndef nextafter_lowering_helper(x, y):","sourceCodeStart":331,"sourceCodeEnd":367,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/utils.py#L331-L367","documentation":"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.","triggerScenarios":"Calling jax.scipy.special.erfinv (or erf_inv) inside a pallas triton kernel on a bfloat16/float16/int array.","commonSituations":"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.","solutions":["Cast to float32 before calling: x.astype(jnp.float32), cast back after","Compute erfinv outside the pallas kernel on the host/XLA side"],"exampleFix":"# before\ny = jax.scipy.special.erfinv(x_bf16)\n# after\ny = jax.scipy.special.erfinv(x_bf16.astype(jnp.float32)).astype(x_bf16.dtype)","handlingStrategy":"validation","validationCode":"assert x.dtype in (jnp.float32, jnp.float64)\nx32 = x.astype(jnp.float32)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Upcast bfloat16/float16 to float32 before special functions in pallas"],"tags":["jax","pallas","triton","dtype","erfinv"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}