Primitive Generic Reduce

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

Usage

prim_reduce(x, init, axes, drop = TRUE, reductor)

Arguments

x: (arrayish)
Arrayish value of any data type.
init: (arrayish)
Scalar (0-dimensional) initial value. Must have the same data type as x and be the neutral element w.r.t. reductor.
axes: (integer())
Axes to reduce over. Negative values count from the end, i.e. -1 refers to the last axis.
drop: (logical(1))
If TRUE (default) the reduced axes are removed; if FALSE they are kept with size 1.
reductor: (function(lhs, rhs))
Binary reducer producing a scalar of the same dtype as x. Must be associative (see "Associativity Requirement").

Value

arrayish
Same data type as x. Shape is x with axes 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 hlo_reduce() with reductor as the body.

Examples

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