Skip to contents

Broadcasts a tensor to a new shape by replicating the data along new or size-1 dimensions.

Usage

nvl_broadcast_in_dim(operand, shape, broadcast_dimensions)

Arguments

operand

(tensorish)
Tensorish value of any data type.

shape

(integer())
Target shape. Each mapped dimension must either match the corresponding operand dimension or the operand dimension must be 1.

broadcast_dimensions

(integer())
Maps each dimension of operand to a dimension of the output. Must have length equal to the number of dimensions of operand.

Value

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

Implemented Rules

  • stablehlo

  • backward

StableHLO

Lowers to stablehlo::hlo_broadcast_in_dim().

Examples

jit_eval({
  x <- nv_tensor(c(1, 2, 3))
  nvl_broadcast_in_dim(x, shape = c(2, 3), broadcast_dimensions = 2L)
})
#> AnvilTensor
#>  1 2 3
#>  1 2 3
#> [ CPUf32{2,3} ]