Primitive Generic Reduce

Reduces an array along the specified dimensions using a user-supplied associative reducer.

Usage

prim_reduce(operand, init, dims, drop = TRUE, reductor)

Arguments

operand: (arrayish)
Arrayish value of any data type.
init: (arrayish)
Scalar (0-dimensional) initial value. Must have the same data type as operand and be the neutral element w.r.t. reductor.
dims: (integer())
Dimensions to reduce over.
drop: (logical(1))
If TRUE (default) the reduced dimensions are removed; if FALSE they are kept with size 1.
reductor: (function(lhs, rhs))
Binary reducer producing a scalar of the same dtype as operand. Must be associative (see "Associativity Requirement").

Value

arrayish
Same data type as operand. Shape is operand with dims removed (or set to 1 if drop = FALSE).

Associativity Requirement

The order in which reductor is applied across the reduction window is implementation-defined. If the reductor is not associative, the result is ill-defined. Furthermore, init must be the neutral element for this reductor. Because floating point math is non-associative, the output of the reduction can differ between backends (GPU, CPU), even if the underlying mathematical function (like +) is associative.

Implemented Rules

stablehlo

StableHLO

Lowers to stablehlo::hlo_reduce() with reductor as the body.

Examples

x <- nv_array(c(1, 2, 3, 4))
prim_reduce(x, init = nv_scalar(0), dims = 1L, reductor = prim_add)
#> AnvlArray
#>  10
#> [ CPUf32{} ] 
prim_reduce(x, init = nv_scalar(1), dims = 1L, reductor = prim_mul)
#> AnvlArray
#>  24
#> [ CPUf32{} ]