jax.lax.broadcast_in_dim

jax.lax.broadcast_in_dim#

jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions)[源代码][源代码]#

封装了 XLA 的 BroadcastInDim 操作符。

参数:

  • operand (ArrayLike) – 一个数组

  • shape (Shape) – 目标数组的形状

  • broadcast_dimensions (Sequence[int]) – 目标形状中的每个维度对应于操作数形状中的哪个维度。也就是说，操作数的第 i 维成为结果的 broadcast_dimensions[i] 维。

返回:

包含结果的数组。

返回类型:

Array

参见

jax.lax.broadcast : 添加新前导维度的更简单接口。

目录