jax.lax.broadcast_in_dim#

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

封装 XLA 的 BroadcastInDim 运算符。

参数:

  • operand (ArrayLike) – 一个数组

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

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

返回:

一个包含结果的数组。

返回类型:

Array

另请参阅

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