jax.numpy.column_stack#

jax.numpy.column_stack(tup)[source]#

Stack arrays column-wise.

JAX implementation of numpy.column_stack().

For arrays of two or more dimensions, this is equivalent to jax.numpy.concatenate() with axis=1.

Parameters:
  • tup (np.ndarray | Array | Sequence[ArrayLike]) – a sequence of arrays to stack; each must have the same leading dimension. Input arrays will be promoted to at least rank 2. If a single array is given it will be treated equivalently to tup = unstack(tup), but the implementation will avoid explicit unstacking.

  • dtype – optional dtype of the resulting array. If not specified, the dtype will be determined via type promotion rules described in Type promotion semantics.

Returns:

the stacked result.

Return type:

Array

See also

Examples

Scalar values:

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

1D arrays:

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

2D arrays:

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