jax.numpy.stack

jax.numpy.stack(arrays, axis=0, out=None, dtype=None)

沿新轴连接数组。

JAX 实现的 numpy.stack()。

参数:

• arrays (np.ndarray | Array | Sequence[ArrayLike]) – 要堆叠的数组序列；每个数组必须具有相同的形状。如果给定的单个数组，它将被视为等同于 arrays = unstack(arrays)，但实现会避免显式地拆堆。

• axis (int) – 指定堆叠的轴。

• out (None | None) – JAX 未使用

• dtype (DTypeLike | None | None) – 结果数组的可选 dtype。如果未指定，则 dtype 将通过 类型提升语义 中描述的类型提升规则确定。

返回:

堆叠后的结果。

返回类型:

Array

另请参阅

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

• jax.numpy.concatenate()：沿现有轴的连接。

• jax.numpy.vstack()：垂直堆叠，即沿轴 0。

• jax.numpy.hstack()：水平堆叠，即沿轴 1。

• jax.numpy.dstack()：深度堆叠，即沿轴 2。

• jax.numpy.column_stack()：堆叠列。

示例

>>> x = jnp.array([1, 2, 3])
>>> y = jnp.array([4, 5, 6])
>>> jnp.stack([x, y])
Array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
>>> jnp.stack([x, y], axis=1)
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)

unstack() 执行逆运算

>>> arr = jnp.stack([x, y], axis=1)
>>> x, y = jnp.unstack(arr, axis=1)
>>> x
Array([1, 2, 3], dtype=int32)
>>> y
Array([4, 5, 6], dtype=int32)