jax-ml/jax · error · NotImplementedError

sign_lowering_helper not implemented for {x.dtype}

Error message

sign_lowering_helper not implemented for {x.dtype}

What it means

pallas.utils.sign_lowering_helper only handles unsigned integers, signed integers, and floating types. Complex or other exotic dtypes raise NotImplementedError since sign is undefined/unsupported in the Triton lowering.

Source

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

  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):
  if x.dtype != y.dtype:
    raise ValueError(
        "The two inputs to `nextafter` must have the same dtype, but got"
        f" {x.dtype} and {y.dtype}"
    )

  if x.dtype not in (jnp.float32, jnp.float64):
    raise ValueError(
        f"`nextafter` only supports float32 and float64, but got {x.dtype}"
    )

  jnp_float, jnp_uint, np_float, np_uint, np_int = (
      jnp.float32, jnp.uint32, np.float32, np.uint32, np.int32,
  ) if x.dtype == jnp.float32 else (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compute sign via real/imag decomposition: sign(x) = sign(x.real) + 1j*sign(x.imag) normalized, or handle complex outside the kernel
2. Cast/reduce to a supported dtype before the call

Example fix

# before
s = jnp.sign(z_complex)  # inside pallas kernel
# after
s = jnp.sign(z.real) + 1j * jnp.sign(z.imag)

if not jnp.issubdtype(x.dtype, jnp.number) or jnp.issubdtype(x.dtype, jnp.complexfloating):
    raise TypeError('sign unsupported in pallas for this dtype')

def sign_supported(x):
    return jnp.issubdtype(x.dtype, jnp.integer) or jnp.issubdtype(x.dtype, jnp.floating)

When it happens

Trigger: Calling jnp.sign on a complex array inside a pallas triton kernel; also boolean inputs depending on subtype checks.

Common situations: Kernels operating on complex64/complex128 that branch through sign; generic numeric code reused across dtypes hitting the Triton path.

Related errors

• Unsupported dtype for sign: {x.dtype}
• unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}
• cannot cast from `{dtype_name}`
• cannot cast to `{dtype_name}`
• Unsigned integer dtype {aval.dtype} is not supported for dot

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