jax-ml/jax · error · NotImplementedError

HLO comparison {direction} for extended dtype {avals_in[0].d

Error message

HLO comparison {direction} for extended dtype {avals_in[0].dtype}

What it means

When lowering a comparison operation on an extended dtype (like jax's ml_dtypes float8 or custom extension types), only EQ and NE comparisons are implemented via special HLO helpers. Requesting any other direction (LT, LE, GT, GE) on an opaque/extended dtype has no lowering, so NotImplementedError is raised.

Source

Thrown at jax/_src/lax/lax.py:5236

  return mlir.delegate_lowering(
      ctx, partial(_unary_reduce_lower, reduction_op, identity,
                   axes=reduce_axes),
      res, avals_in=[base_aval_out], avals_out=[aval_out])

_opaque_eq_hlo = partial(
    _opaque_comparison_hlo, 'EQ', hlo.AndOp, _get_bitwise_and_identity)
_opaque_ne_hlo = partial(
    _opaque_comparison_hlo, 'NE', hlo.OrOp, _get_bitwise_or_identity)

def _compare_lower_hlo_opaque(direction: str, ctx, avals_in, aval_out, x, y):
  broadcast_avals_in = tuple(
      core.ShapedArray(aval_out.shape, aval.dtype) for aval in avals_in)
  if direction == 'EQ':
    return _opaque_eq_hlo(ctx, broadcast_avals_in, aval_out, x, y)
  elif direction == 'NE':
    return _opaque_ne_hlo(ctx, broadcast_avals_in, aval_out, x, y)
  else:
    raise NotImplementedError(
        f"HLO comparison {direction} for extended dtype {avals_in[0].dtype}")


def _compare_lower_hlo(direction: str, total_order: bool, ctx, x, y):
  avals_in, (aval_out,) = ctx.avals_in, ctx.avals_out
  x_dtype = avals_in[0].dtype
  x, y = mlir.multi_broadcast_in_dim(ctx, (x, y), avals_in, aval_out.shape,
                                     aval_out.sharding)
  if dtypes.issubdtype(x_dtype, dtypes.extended):
    assert not total_order
    return _compare_lower_hlo_opaque(direction, ctx, avals_in, aval_out, x, y)
  if dtypes.issubdtype(x_dtype, np.inexact):
    compare_type = "TOTALORDER" if total_order else "FLOAT"
  elif dtypes.issubdtype(x_dtype, np.signedinteger):
    compare_type = "SIGNED"
  else:
    compare_type = "UNSIGNED"
  return [mlir.compare_hlo(x, y, direction, compare_type)]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use equality comparisons (jnp.equal / jnp.not_equal) which are supported
Convert to a standard dtype (e.g. .astype(jnp.float32)) before ordering comparisons
For float8, compare via bitcast to uint8 only if you understand the bit layout
Implement/extend the dtype rules if it is a custom extended dtype

Example fix

// before
mask = x_f8 < y_f8

// after
mask = x_f8.astype(jnp.float32) < y_f8.astype(jnp.float32)

Defensive patterns

Strategy: type-guard

Validate before calling

def is_orderable(dtype):
    return not isinstance(dtype, jax.dtypes.ExtendedDType)

Type guard

def orderable(x):
    return not isinstance(x.dtype, jax.dtypes.ExtendedDType)

Try / catch

try:
    mask = x < y
except NotImplementedError:
    mask = x.astype(jnp.float32) < y.astype(jnp.float32)

Prevention

Reserve ordering comparisons for standard dtypes
Convert extended dtypes to float32 before sort/min/max/mask logic
Document which ops are EQ/NE-only for extended dtypes

When it happens

Trigger: Calling jnp.less/greater/less_equal/greater_equal (or lax.lt etc.) on arrays with an ExtendedDType whose comparison is opaque, e.g. comparing float8_* or custom extension dtypes; equality (==, !=) works but ordering does not.

Common situations: Using float8 dtypes with jnp.sort, jnp.maximum, clipping, or boolean masks that lower to ordered comparisons; assuming all numpy comparison semantics carry over to new extended dtypes.

Related errors

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