jax.numpy.unstack

jax.numpy.unstack(x, /, *, axis=0)

沿轴展开数组。

array_api.unstack() 的 JAX 实现。

参数:

• x (ArrayLike) – 要展开的数组。必须有 x.ndim >= 1。

• axis (int) – 沿其展开的整数轴。必须满足 -x.ndim <= axis < x.ndim。

返回:

展开数组的元组。

返回类型:

tuple[Array, …]

另请参阅

• jax.numpy.stack()：unstack 的逆运算

• jax.numpy.split()：沿轴将数组拆分为批次。

示例

>>> arr = jnp.array([[1, 2, 3],
...                  [4, 5, 6]])
>>> arrs = jnp.unstack(arr)
>>> print(*arrs)
[1 2 3] [4 5 6]

stack() 提供了这个的逆运算

>>> jnp.stack(arrs)
Array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)