jax.numpy.broadcast_to

jax.numpy.broadcast_to(array, shape)

将数组广播到指定的形状。

numpy.broadcast_to() 的 JAX 实现。JAX 使用 NumPy 风格的广播规则，您可以在 NumPy 广播 中了解更多相关信息。

参数:

• array (ArrayLike) – 要广播的数组。

• shape (DimSize | Shape) – 数组将广播到的形状。

返回值:

广播到指定形状的数组副本。

返回类型:

Array

参见

• jax.numpy.broadcast_arrays(): 将数组广播到公共形状。

• jax.numpy.broadcast_shapes(): 将输入形状广播到公共形状。

示例

>>> x = jnp.int32(1)
>>> jnp.broadcast_to(x, (1, 4))
Array([[1, 1, 1, 1]], dtype=int32)

>>> x = jnp.array([1, 2, 3])
>>> jnp.broadcast_to(x, (2, 3))
Array([[1, 2, 3],
       [1, 2, 3]], dtype=int32)

>>> x = jnp.array([[2], [4]])
>>> jnp.broadcast_to(x, (2, 4))
Array([[2, 2, 2, 2],
       [4, 4, 4, 4]], dtype=int32)