{"record":{"id":"611f3d5d69b8f4d9","repo":"xai-org/x-algorithm","slug":"only-int32-is-supported-for-unique","errorCode":null,"errorMessage":"Only int32 is supported for unique.","messagePattern":"Only int32 is supported for unique\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/unique/__init__.py","lineNumber":18,"sourceCode":"# SPDX-License-Identifier: Apache-2.0\n# Copyright 2026 X.AI Corp.\nimport jax\nimport jax.numpy as jnp\n\ntry:\n    from xrex.cuda.unique.src import unique_api\nexcept ImportError:\n    unique_api = None\nelse:\n    jax.ffi.register_ffi_target(\"xrex_unique\", fn=unique_api.unique(), platform=\"CUDA\")\n\n\ndef unique(\n    x: jax.Array, return_inverse: bool, size: int, fill_value: int\n) -> tuple[jax.Array, jax.Array]:\n    if x.dtype != jnp.int32:\n        raise ValueError(\"Only int32 is supported for unique.\")\n    if x.ndim != 1:\n        raise ValueError(\"Only 1D arrays are supported for unique.\")\n\n    if unique_api is None or jax.default_backend() != \"gpu\":\n        unique_vals, unique_inverse = jnp.unique(\n            x, return_inverse=return_inverse, size=size, fill_value=fill_value\n        )\n        return unique_vals.astype(x.dtype), unique_inverse.astype(x.dtype)\n\n    call = jax.ffi.ffi_call(\n        \"xrex_unique\",\n        [\n            jax.ShapeDtypeStruct(shape=[size], dtype=x.dtype),\n            jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),\n            jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),\n            jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),\n            jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),\n        ],","sourceCodeStart":1,"sourceCodeEnd":36,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/unique/__init__.py#L1-L36","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Cast the input before calling: x = x.astype(jnp.int32)","Ensure upstream index-generating ops (argsort, nonzero, searchsorted) produce int32","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)"],"exampleFix":"// before\nvals, inv = unique(user_ids, return_inverse=True, size=n, fill_value=-1)  # user_ids is int64\n// after\nvals, inv = unique(user_ids.astype(jnp.int32), return_inverse=True, size=n, fill_value=-1)","handlingStrategy":"type-guard","validationCode":"if x.dtype != jnp.int32:\n    x = x.astype(jnp.int32)\nvals, inv = unique(x, return_inverse=True, size=size, fill_value=fill_value)","typeGuard":"def is_int32_1d(x: jax.Array) -> bool:\n    return x.dtype == jnp.int32 and x.ndim == 1","tryCatchPattern":null,"preventionTips":["Standardize id/feature-unique pipelines on int32 from ingestion","Check x.dtype after numpy interop (numpy defaults to int64 on Linux)","Centralize the cast in wrapper helpers that call unique"],"tags":["cuda","jax","unique","dtype-validation"],"backgroundTag":"unsupported-dtype","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}