jax.numpy.concat — JAX documentation (original) (raw)

jax.numpy.concat#

jax.numpy.concat(arrays, /, *, axis=0)[source]#

Join arrays along an existing axis.

JAX implementation of array_api.concat().

Parameters:

• arrays (Sequence [ ArrayLike ]) – a sequence of arrays to concatenate; each must have the same shape except along the specified axis. If a single array is given it will be treated equivalently to arrays = unstack(arrays), but the implementation will avoid explicit unstacking.
• axis (int | None) – specify the axis along which to concatenate.

Returns:

the concatenated result.

Return type:

Array

See also

• jax.lax.concatenate(): XLA concatenation API.
• jax.numpy.concatenate(): NumPy version of this function.
• jax.numpy.stack(): concatenate arrays along a new axis.

Examples

One-dimensional concatenation:

x = jnp.arange(3) y = jnp.zeros(3, dtype=int) jnp.concat([x, y]) Array([0, 1, 2, 0, 0, 0], dtype=int32)

Two-dimensional concatenation:

x = jnp.ones((2, 3)) y = jnp.zeros((2, 1)) jnp.concat([x, y], axis=1) Array([[1., 1., 1., 0.], [1., 1., 1., 0.]], dtype=float32)