JAX：JIT编译函数中计算数组唯一元素数量的替代方案咨询

由于JAX JIT要求张量形状在编译时确定，jax.numpy.unique因返回动态尺寸结果会触发ConcretizationTypeError，以下是几种算法层面的可行解决办法：

1. 排序后统计相邻差异 #

思路 #

将数组排序后，相同元素会连续排列。通过比较相邻元素的差异，统计不同元素的数量（差异不为0的次数+1，因为第一个元素默认算一个）。所有操作的张量形状均为静态，符合JIT要求。

代码示例 #

import jax.numpy as jnp

def count_unique_sorted(x):
    if x.size == 0:
        return 0
    sorted_x = jnp.sort(x)
    # 前置一个不可能出现在数组中的值，确保第一个元素被计数
    diffs = jnp.diff(sorted_x, prepend=-jnp.inf)
    return jnp.sum(diffs != 0)

适用场景 #

适用于任意可排序的元素类型（整数、浮点数等），时间复杂度为O(n log n)，是通用程度较高的方案。

2. 固定范围的计数桶（针对整数元素） #

思路 #

若已知数组元素的取值范围，可使用jnp.bincount统计每个值的出现次数，再统计次数大于0的桶的数量。所有操作的输出形状均为静态。

代码示例 #

def count_unique_bincount(x, min_val=0, max_val=100):
    # 确保覆盖所有可能的元素值
    counts = jnp.bincount(x, minlength=max_val - min_val + 1)
    return jnp.sum(counts > 0)

适用场景 #

仅适用于整数元素且取值范围已知的场景，时间复杂度为O(n)，是效率最高的方案。若为浮点数，可先通过缩放+取整转化为整数（如jnp.round(x * 100).astype(int)）后使用此方法。

3. 标记元素首次出现的位置 #

思路 #

对数组中的每个元素，判断它在当前位置之前是否已出现过，统计首次出现的元素数量。通过广播比较实现，所有中间张量形状均为静态。

代码示例 #

def count_unique_first_occurrence(x):
    n = x.size
    if n == 0:
        return 0
    # 创建索引矩阵，标记当前位置之前的元素
    idx_matrix = jnp.arange(n)[:, None] > jnp.arange(n)[None, :]
    # 比较当前元素与之前所有元素是否相等
    has_occurred = jnp.any(idx_matrix & (x[:, None] == x[None, :]), axis=1)
    # 统计未出现过的元素数量（即首次出现的元素）
    return jnp.sum(~has_occurred)

适用场景 #

适用于小尺寸数组（n较小），时间复杂度为O(n²)，n较大时效率较低。

4. 固定容量的唯一元素缓冲区（已知最大唯一数上限） #

思路 #

若预先知道数组中唯一元素的最大数量K，可初始化一个固定大小的缓冲区，遍历数组时将未在缓冲区中的元素加入，最后统计缓冲区中有效元素的数量。使用jax.lax.fori_loop实现循环，符合JIT要求。

代码示例 #

from jax import lax

def count_unique_fixed_buffer(x, max_unique=10):
    n = x.size
    if n == 0:
        return 0
    # 初始化缓冲区，用不可能出现的值填充
    init_buf = jnp.full((max_unique,), -jnp.inf)
    
    def loop_body(i, state):
        buf, count = state
        val = x[i]
        # 检查当前值是否已在缓冲区的有效部分
        exists = jnp.any(buf[:count] == val)
        # 更新缓冲区和计数
        new_count = lax.cond(exists, lambda c: c, lambda c: c + 1, count)
        new_buf = lax.cond(exists, lambda b: b, lambda b: b.at[count].set(val), buf)
        return (new_buf, new_count)
    
    _, final_count = lax.fori_loop(0, n, loop_body, (init_buf, 0))
    return final_count

适用场景 #

适用于已知唯一元素最大数量的场景，时间复杂度为O(n*K)，K较小时效率较高。

内容的提问来源于stack exchange，提问作者blindeyes