xai-org/x-algorithm · error · ValueError

Only int32 is supported for unique.

Error message

Only int32 is supported for unique.

What it means

The custom CUDA-accelerated unique kernel in phoenix/xrex/cuda/unique only implements the int32 dtype path; any other dtype is rejected up front with a ValueError instead of producing silently wrong results. On non-GPU backends or when the API is unavailable it falls back to jnp.unique, but the dtype check applies regardless.

Source

Thrown at phoenix/xrex/cuda/unique/__init__.py:18

# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 X.AI Corp.
import jax
import jax.numpy as jnp

try:
    from xrex.cuda.unique.src import unique_api
except ImportError:
    unique_api = None
else:
    jax.ffi.register_ffi_target("xrex_unique", fn=unique_api.unique(), platform="CUDA")


def unique(
    x: jax.Array, return_inverse: bool, size: int, fill_value: int
) -> tuple[jax.Array, jax.Array]:
    if x.dtype != jnp.int32:
        raise ValueError("Only int32 is supported for unique.")
    if x.ndim != 1:
        raise ValueError("Only 1D arrays are supported for unique.")

    if unique_api is None or jax.default_backend() != "gpu":
        unique_vals, unique_inverse = jnp.unique(
            x, return_inverse=return_inverse, size=size, fill_value=fill_value
        )
        return unique_vals.astype(x.dtype), unique_inverse.astype(x.dtype)

    call = jax.ffi.ffi_call(
        "xrex_unique",
        [
            jax.ShapeDtypeStruct(shape=[size], dtype=x.dtype),
            jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),
            jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),
            jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),
            jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),
        ],

View on GitHub (pinned to 24c60942c5)

Solutions

1. Cast the input before calling: x = x.astype(jnp.int32)
2. Ensure upstream index-generating ops (argsort, nonzero, searchsorted) produce int32
3. If values can exceed int32 range, remap/compress ids to a dense int32 range first (e.g. via this unique itself or an encoding step)

Example fix

// before
vals, inv = unique(user_ids, return_inverse=True, size=n, fill_value=-1)  # user_ids is int64
// after
vals, inv = unique(user_ids.astype(jnp.int32), return_inverse=True, size=n, fill_value=-1)

if x.dtype != jnp.int32:
    x = x.astype(jnp.int32)
vals, inv = unique(x, return_inverse=True, size=size, fill_value=fill_value)

def is_int32_1d(x: jax.Array) -> bool:
    return x.dtype == jnp.int32 and x.ndim == 1

When it happens

Trigger: Calling unique(x, ...) with x of dtype int64, uint32, float32, bfloat16, etc., from any caller such as prefetcher_loop, _sample_topic_pool, or sample_engaged_posts.

Common situations: Feeding indices produced by jnp.argsort or argpartition (often int32) after an astype to int64; loading data with numpy defaults (int64 on Linux) and passing it through; JAX default integer promotion changing dtype between versions.

Related errors

• Only 1D arrays are supported for unique.
• ASTNode %s expected %d arguments, %d passed.
• k ({k}) must be <= n ({n})
• Only Float16 or BFloat16 is supported
• q, k, and v should all have the same dtype, got: {q.dtype},

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/611f3d5d69b8f4d9. Report an issue: GitHub.