jax.numpy.unique

目录

jax.numpy.unique#

jax.numpy.unique(ar, return_index=False, return_inverse=False, return_counts=False, axis=None, *, equal_nan=True, size=None, fill_value=None)[源代码][源代码]#

返回数组中的唯一值。

JAX 实现的 numpy.unique()。

因为 unique 的输出大小取决于数据，该函数通常不与 jit() 和其他 JAX 变换兼容。JAX 版本添加了可选的 size 参数，必须在 jnp.unique 用于此类上下文中时静态指定该参数。

参数:

  • ar (ArrayLike) – 将从中提取唯一值的 N 维数组。

  • return_index (bool) – 如果为真，还返回 ar 中每个值出现的索引

  • return_inverse (bool) – 如果为真，还返回可以用于从唯一值重建 ar 的索引。

  • return_counts (bool) – 如果为真，还返回每个唯一值的出现次数。

  • axis (int | None) – 如果指定，则沿指定轴计算唯一值。如果为 None（默认），则在计算唯一值之前先展平 ar。

  • equal_nan (bool) – 如果为 True，则在确定唯一性时将 NaN 值视为等价。

  • size (int | None) – 如果指定，则仅返回前 size 个排序后的唯一元素。如果唯一元素的数量少于 size 指示的数量，则返回值将用 fill_value 填充。

  • fill_value (ArrayLike | None) – 当指定了 size 且元素数量少于指定数量时，用 fill_value 填充剩余的条目。默认为最小唯一值。

返回:

一个数组或数组的元组，取决于 return_index、return_inverse 和 return_counts 的值。返回的值是 - unique_values: 如果 axis 是 None，一个长度为 n_unique 的 1D 数组。如果 axis 被指定，形状为 (*ar.shape[:axis], n_unique, *ar.shape[axis + 1:])。 - unique_index: (仅在 return_index 为 True 时返回) 形状为 (n_unique,) 的数组。包含 ar 中每个唯一值首次出现的索引。对于 1D 输入，ar[unique_index] 等同于 unique_values。 - unique_inverse: (仅在 return_inverse 为 True 时返回) 如果 axis 是 None，形状为 (ar.size,) 的数组，如果 axis 被指定，形状为 (ar.shape[axis],)。包含 ar 中每个值在 unique_values 中的索引。对于 1D 输入，unique_values[unique_inverse] 等同于 ar。 - unique_counts: (仅在 return_counts 为 True 时返回) 形状为 (n_unique,) 的数组。包含 ar 中每个唯一值的出现次数。

参见

示例

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> jnp.unique(x)
Array([1, 3, 4], dtype=int32)

即时编译与大小参数

如果你在 jit() 或其他变换下尝试这个，你会得到一个错误，因为输出形状是动态的：

>>> jax.jit(jnp.unique)(x)  
Traceback (most recent call last):
   ...
jax.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[5].
The error arose for the first argument of jnp.unique(). To make jnp.unique() compatible with JIT and other transforms, you can specify a concrete value for the size argument, which will determine the output size.

问题在于，转换后的函数的输出必须具有静态形状。为了实现这一点，你可以传递一个静态的 size 参数：

>>> jit_unique = jax.jit(jnp.unique, static_argnames=['size'])
>>> jit_unique(x, size=3)
Array([1, 3, 4], dtype=int32)

如果你的静态大小小于实际的唯一值数量，它们将被截断。

>>> jit_unique(x, size=2)
Array([1, 3], dtype=int32)

如果静态大小大于实际唯一值的数量，它们将被填充为 fill_value，默认值为最小唯一值：

>>> jit_unique(x, size=5)
Array([1, 3, 4, 1, 1], dtype=int32)
>>> jit_unique(x, size=5, fill_value=0)
Array([1, 3, 4, 0, 0], dtype=int32)

多维唯一值

如果你将一个多维数组传递给 unique，它将默认被展平：

>>> M = jnp.array([[1, 2],
...                [2, 3],
...                [1, 2]])
>>> jnp.unique(M)
Array([1, 2, 3], dtype=int32)

如果你传递一个 axis 关键字，你可以沿着该轴找到数组的唯一 切片：

>>> jnp.unique(M, axis=0)
Array([[1, 2],
       [2, 3]], dtype=int32)

返回索引

如果你设置 return_index=True，那么 unique 返回每个唯一值首次出现的索引：

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> values, indices = jnp.unique(x, return_index=True)
>>> print(values)
[1 3 4]
>>> print(indices)
[2 0 1]
>>> jnp.all(values == x[indices])
Array(True, dtype=bool)

在多维情况下，可以使用沿指定轴计算的 jax.numpy.take() 提取唯一值：

>>> values, indices = jnp.unique(M, axis=0, return_index=True)
>>> jnp.all(values == jnp.take(M, indices, axis=0))
Array(True, dtype=bool)

返回逆矩阵

如果你设置 return_inverse=True ，那么 unique 会返回输入数组中每个条目在唯一值中的索引：

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> values, inverse = jnp.unique(x, return_inverse=True)
>>> print(values)
[1 3 4]
>>> print(inverse)
[1 2 0 1 0]
>>> jnp.all(values[inverse] == x)
Array(True, dtype=bool)

在多维情况下，可以使用 jax.numpy.take() 来重建输入：

>>> values, inverse = jnp.unique(M, axis=0, return_inverse=True)
>>> jnp.all(jnp.take(values, inverse, axis=0) == M)
Array(True, dtype=bool)

返回计数

如果你设置 return_counts=True ，那么 unique 会返回输入中每个唯一值的出现次数：

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> values, counts = jnp.unique(x, return_counts=True)
>>> print(values)
[1 3 4]
>>> print(counts)
[2 2 1]

对于多维数组，这也返回一个一维的计数数组，指示沿指定轴的出现次数：

>>> values, counts = jnp.unique(M, axis=0, return_counts=True)
>>> print(values)
[[1 2]
 [2 3]]
>>> print(counts)
[2 1]
目录