jax.numpy.concatenate

内容

jax.numpy.concatenate#

jax.numpy.concatenate(arrays, axis=0, dtype=None)[source]#

沿着现有轴连接数组。

JAX 实现 numpy.concatenate().

参数:
  • arrays (np.ndarray | Array | Sequence[ArrayLike]) – 要连接的数组序列;除了指定轴之外,每个数组都必须具有相同的形状。如果给出一个单一的数组,它将等同于 arrays = unstack(arrays),但实现将避免显式拆解。

  • axis (int | None) – 指定要连接的轴。

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

返回值:

连接的结果。

返回值类型:

Array

参见

示例

一维连接

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

二维连接

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