Stack arrays depth-wise.
JAX implementation of numpy.dstack().
For arrays of three or more dimensions, this is equivalent to jax.numpy.concatenate() with axis=2.
tup (np.ndarray | Array | Sequence[ArrayLike]) – a sequence of arrays to stack; each must have the same shape along all but the third axis. Input arrays will be promoted to at least rank 3. If a single array is given it will be treated equivalently to tup = unstack(tup), but the implementation will avoid explicit unstacking.
dtype (DTypeLike | None) – optional dtype of the resulting array. If not specified, the dtype will be determined via type promotion rules described in Type promotion semantics.
the stacked result.
Examples
Scalar values:
>>> jnp.dstack([1, 2, 3]) Array([[[1, 2, 3]]], dtype=int32, weak_type=True)
1D arrays:
>>> x = jnp.arange(3)
>>> y = jnp.ones(3)
>>> jnp.dstack([x, y])
Array([[[0., 1.],
[1., 1.],
[2., 1.]]], dtype=float32)
2D arrays:
>>> x = x.reshape(1, 3)
>>> y = y.reshape(1, 3)
>>> jnp.dstack([x, y])
Array([[[0., 1.],
[1., 1.],
[2., 1.]]], dtype=float32)
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4