{"record":{"id":"e9177c8767f962a2","repo":"xai-org/x-algorithm","slug":"only-1d-arrays-are-supported-for-unique","errorCode":null,"errorMessage":"Only 1D arrays are supported for unique.","messagePattern":"Only 1D arrays are supported for unique\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/unique/__init__.py","lineNumber":20,"sourceCode":"# 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        ],\n    )\n    unique_vals, unique_inverse, _, _, _ = call(","sourceCodeStart":2,"sourceCodeEnd":38,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/unique/__init__.py#L2-L38","documentation":"The custom unique kernel only supports 1D arrays, matching the semantics of a single flat unique operation. Multi-dimensional inputs are rejected with a ValueError before the GPU path (or the jnp.unique fallback) is taken, because the kernel cannot iterate higher-rank layouts.","triggerScenarios":"Calling unique(x, ...) where x.ndim != 1, e.g. a (batch, seq) tensor of ids from prefetcher_loop or sample_engaged_posts.","commonSituations":"Switching from per-row unique to batched tensors without reshaping; forgetting that a (1, N) shaped slice still has ndim == 2; passing token-id matrices instead of flattened id lists.","solutions":["Flatten before calling: unique(x.reshape(-1), ...) and reshape the inverse output back afterward","If per-row uniqueness is needed, vmap or loop over rows instead of passing 2D","Check x.ndim and x.shape immediately after data loading to catch this early"],"exampleFix":"// before\nvals, inv = unique(ids_2d, return_inverse=True, size=size, fill_value=-1)\n// after\nvals, inv = unique(ids_2d.reshape(-1), return_inverse=True, size=size, fill_value=-1)\ninv = inv.reshape(ids_2d.shape)","handlingStrategy":"validation","validationCode":"if x.ndim != 1:\n    orig_shape = x.shape\n    x = x.reshape(-1)\nvals, inv = unique(x, ...)\n# inv = inv.reshape(orig_shape) if needed","typeGuard":"def is_flat_int32(x: jax.Array) -> bool:\n    return x.ndim == 1 and x.dtype == jnp.int32","tryCatchPattern":null,"preventionTips":["Reshape at data-loading boundaries so downstream kernels always see 1D","Use jax.vmap for per-row uniqueness rather than passing 2D tensors"],"tags":["cuda","jax","unique","shape-validation"],"backgroundTag":"invalid-input-shape","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}