jax-ml/jax · error · ValueError

Unsupported dtype for sign: {x.dtype}

Error message

Unsupported dtype for sign: {x.dtype}

What it means

The Mosaic sign lowering handles floats, signed ints (via comparison trick), and unsigned ints. Any other dtype (e.g. complex, bool) raises ValueError('Unsupported dtype for sign: ...'). Note this is a ValueError, not NotImplementedError, and comes from the Python-level _lower_fun helper.

Source

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

      sign_val = lax.bitcast_convert_type(sign_val_i, jnp.float32)
      # By checking abs(x32) > 0.0 we handle NaN and +/-0.0.
      res = jnp.where(jnp.abs(x32) > 0.0, sign_val, x32)

      if dtype == jnp.bfloat16:
        assert not tpu_has_native_bf16
        # Drop the rightmost 16 bits, which are all zero.
        res_i = lax.bitcast_convert_type(res, jnp.uint32)
        res_u16 = lax.convert_element_type(
            lax.shift_right_logical(res_i, jnp.uint32(16)), jnp.uint16
        )
        return lax.bitcast_convert_type(res_u16, jnp.bfloat16)
      return res.astype(dtype)

    if jnp.issubdtype(x.dtype, jnp.signedinteger):
      return (x > 0).astype(x.dtype) - (x < 0).astype(x.dtype)
    if jnp.issubdtype(x.dtype, jnp.unsignedinteger):
      return (x != 0).astype(x.dtype)
    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")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Define sign semantics yourself for complex (e.g. z/|z|) and implement manually
2. Cast to float32/int32 before sign
3. Avoid jnp.sign on bools — use the mask directly

Example fix

// before
s = jnp.sign(z)  # complex
// after
mag = jnp.sqrt(z.real**2 + z.imag**2)
s = jnp.where(mag > 0, z / mag, 0)

import jax.numpy as jnp
def sign_dtype_ok(dt):
    return (jnp.issubdtype(dt, jnp.floating) or jnp.issubdtype(dt, jnp.signedinteger)
            or jnp.issubdtype(dt, jnp.unsignedinteger))

def is_sign_supported(dt): return not jnp.issubdtype(dt, jnp.complexfloating)

When it happens

Trigger: Calling jnp.sign/lax.sign on complex or boolean arrays inside a Pallas kernel.

Common situations: Sign of complex numbers (ill-defined anyway); sign of booleans; unexpected dtype promotion.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• 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/80e7d6974fdc7bee. Report an issue: GitHub.