jax.numpy.dstack

jax.numpy.dstack(tup, dtype=None)

沿深度方向堆叠数组。

JAX 对 numpy.dstack() 的实现。

对于三维或更高维度的数组，这等效于使用 axis=2 的 jax.numpy.concatenate()。

参数：

• tup (np.ndarray | Array | Sequence[ArrayLike]) – 要堆叠的数组序列；除了第三轴之外，每个数组必须具有相同的形状。输入数组将至少提升到秩 3。如果给定的只是一个数组，则它将被视为等效于 tup = unstack(tup)，但实现将避免显式拆堆。

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

返回：

堆叠的结果。

返回类型：

Array

另请参阅

• jax.numpy.stack()：沿任意轴堆叠

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

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

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

示例

标量值

>>> jnp.dstack([1, 2, 3])
Array([[[1, 2, 3]]], dtype=int32, weak_type=True)

一维数组

>>> x = jnp.arange(3)
>>> y = jnp.ones(3)
>>> jnp.dstack([x, y])
Array([[[0., 1.],
        [1., 1.],
        [2., 1.]]], dtype=float32)

二维数组

>>> x = x.reshape(1, 3)
>>> y = y.reshape(1, 3)
>>> jnp.dstack([x, y])
Array([[[0., 1.],
        [1., 1.],
        [2., 1.]]], dtype=float32)