jax.lax.broadcast_in_dim#

jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions, *, out_sharding=None)[source]#

包装了 XLA 的 BroadcastInDim 运算符。

参数:

  • operand (ArrayLike) – 一个数组

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

  • broadcast_dimensions (Sequence[int]) – operand 形状的维度 i 对应于目标形状的维度 broadcast_dimensions[i]。

返回:

包含结果的数组。

返回类型:

Array

另请参阅

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