jax.numpy.concat#

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

Join arrays along an existing axis.

JAX implementation ofarray_api.concat().

Parameters:

• arrays (Sequence[ArrayLike]) – a sequence of arrays to concatenate; each must have the same shapeexcept along the specified axis. If a single array is given it will betreated equivalently toarrays = unstack(arrays), but the implementationwill 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)