jax.numpy.broadcast_arrays

jax.numpy.broadcast_arrays(*args)

将数组广播到通用形状。

numpy.broadcast_arrays() 的 JAX 实现。JAX 使用 NumPy 风格的广播规则，您可以在NumPy 广播中阅读更多内容。

参数:

args (ArrayLike) – 零个或多个要进行广播的类数组对象。

返回值:

包含输入广播副本的数组列表。

返回类型:

list[Array]

另请参阅

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

• jax.numpy.broadcast_to(): 将数组广播为指定形状。

示例

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

>>> x = jnp.array([[1, 2, 3]])
>>> y = jnp.array([[10],
...                [20]])
>>> x2, y2 = jnp.broadcast_arrays(x, y)
>>> x2
Array([[1, 2, 3],
       [1, 2, 3]], dtype=int32)
>>> y2
Array([[10, 10, 10],
       [20, 20, 20]], dtype=int32)