jax.numpy.unstack

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

沿一个轴解开数组。

JAX 实现 array_api.unstack()。

参数:

• 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)