jax-ml/jax · error · ValueError

`nextafter` only supports float32 and float64, but got {x.dt

Error message

`nextafter` only supports float32 and float64, but got {x.dtype}

What it means

The Pallas Triton lowering of nextafter only implements the bit-twiddling algorithm for float32 and float64. Other dtypes (float16, bfloat16, integers) raise ValueError before lowering.

Source

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

    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 (
      jnp.float64, jnp.uint64, np.float64, np.uint64, np.int64,
  )

  bitwidth = dtypes.itemsize_bits(x.dtype)

  x_as_int = x.view(jnp_uint)
  y_as_int = y.view(jnp_uint)

  # The result is NaN if either "x" or "y" are NaN.
  nan_input = jnp.isnan(x) | jnp.isnan(y)
  result_for_nan = jnp.full_like(x_as_int, np_float(np.nan).view(np_uint))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast to float32: jnp.nextafter(x.astype(jnp.float32), y.astype(jnp.float32)).astype(x.dtype)
Move the nextafter computation outside the pallas kernel

Example fix

# before
jnp.nextafter(x_bf16, y_bf16)
# after
jnp.nextafter(x_bf16.astype(jnp.float32), y_bf16.astype(jnp.float32)).astype(jnp.bfloat16)

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

Compute nextafter in float32 inside pallas kernels

When it happens

Trigger: Calling jnp.nextafter on bfloat16 or float16 tensors inside a pallas triton kernel.

Common situations: Half-precision kernels on modern accelerators where bfloat16 is the default dtype calling nextafter (e.g. for epsilon stepping in numerics).

Related errors

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