Skip to contents

Reduces an array along the specified axes using a user-supplied associative reducer. reducer and init must satisfy the constraints in the "Associativity Requirement" section below.

Usage

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

Arguments

x

(arrayish)
The array to reduce. Can be any data type. x and init must have the same data type. An R value among them assumes the data type of the others when it is in its data type category, and its default data type when none of them has one.

init

(arrayish)
Scalar initial value, with no axes. Must be the neutral element w.r.t. reducer, and shares x's data type – see x.

axes

(integer())
Axes to reduce over. Negative values count from the end, i.e. -1 refers to the last axis.

reducer

(function(lhs, rhs))
Binary reducer producing a scalar of the same data type as x. Its two arguments are passed by position, so they may carry any names. Must be associative (see "Associativity Requirement").

drop

(logical(1))
If TRUE (default) the reduced axes are removed; if FALSE they are kept with size 1.

Value

(arrayish)
Has the input's data type. The shape is the input's with the reduced axes removed (drop = TRUE) or set to 1 (drop = FALSE).

Associativity Requirement

The order in which reducer is applied across the reduction window is implementation-defined. If the reducer is not associative, the result is ill-defined. Furthermore, init must be the neutral element for this reducer. Because arithmetic in a float data type 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(), specified under reduce. The body is reducer.

Examples

# `init` shares `x`'s data type, and the reduced axis disappears
x <- nv_array(c(1, 2, 3, 4))
prim_reduce(x, init = nv_scalar(0), axes = 1L, reducer = prim_add)
#> AnvlArray
#>  10
#> [ CPUf32{} ] 
prim_reduce(x, init = nv_scalar(1), axes = 1L, reducer = prim_mul)
#> AnvlArray
#>  24
#> [ CPUf32{} ]