jax.lax.broadcast_in_dim#

jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions, sharding=None)[原始碼]#

包裝 XLA 的 BroadcastInDim 運算子。

參數:

  • operand (ArrayLike) – 一個陣列

  • shape (Shape) – 目標陣列的形狀

  • broadcast_dimensions (Sequence[int]) – 運算元形狀的每個維度對應到目標形狀中的哪個維度。也就是說，運算元的維度 i 變成結果的維度 broadcast_dimensions[i]。

回傳:

包含結果的陣列。

回傳型別:

Array

另請參閱

jax.lax.broadcast：用於新增前導維度的更簡單介面。