Skip to contents

Broadcasts an array to a new shape by replicating the data along new or size-1 axes.

Usage

prim_broadcast_in_axes(x, shape, broadcast_axes)

Arguments

x

(arrayish)
Arrayish value of any data type.

shape

(integer())
Target shape. The size of each mapped axis must either match the size of the corresponding axis of x, or that axis of x must have size 1.

broadcast_axes

(integer())
Maps each axis of x to an axis of the output. Must have length equal to the number of axes of x.

Value

arrayish
Has the same data type as the input and the given shape. It is ambiguous if the input is ambiguous.

Implemented Rules

  • stablehlo

  • quickr

  • reverse

StableHLO

Lowers to hlo_broadcast_in_dim().

Examples

x <- nv_array(c(1, 2, 3))
prim_broadcast_in_axes(x, shape = c(2, 3), broadcast_axes = 2L)
#> AnvlArray
#>  1 2 3
#>  1 2 3
#> [ CPUf32{2,3} ]