Primitive Scan

Runs body a fixed number of times, threading a carry through the steps and stacking each step's outputs along a new leading axis. Step t receives the carry and, for every array in xs, its slice at position t along axis 1 with that axis dropped.

Usage

prim_scan(init, xs, body, length, reverse = FALSE)

Arguments

init: (list())
Initial carry: a (possibly nested) list of arrays. Every leaf keeps its shape and data type across steps.
xs: (list())
Per-step inputs: a (possibly nested) list of arrays sliced along axis 1, all of size length along it. An empty list runs a counted loop.
body: (function)
Step function function(carry, x) returning list(carry = , out = ), where carry has the structure of init and out is a (possibly nested) list of arrays or NULL. x is NULL when xs is empty.
length: (integer(1))
Static trip count; the size of axis 1 of every array in xs. 0 runs no step and returns init with zero-length stacked outputs.
reverse: (logical(1))
If TRUE, steps run from length down to 1; each step still reads xs at its own position and writes its output there.

Value

list(carry = , out = ): the final carry and the stacked outputs, each leaf of out gaining a leading axis of size length.

Implemented Rules

stablehlo

StableHLO

Lowers to hlo_while() over a counter, the carry, the output buffers and xs, with hlo_dynamic_slice() reading each step's inputs and hlo_dynamic_update_slice() writing its outputs.

Examples

prim_scan(
  init = list(s = nv_scalar(0)),
  xs = list(x = nv_array(c(1, 2, 3))),
  body = function(carry, x) {
    s <- carry$s + x$x
    list(carry = list(s = s), out = s)
  },
  length = 3L
)
#> $carry
#> $carry$s
#> AnvlArray
#>  6
#> [ CPUf32{} ] 
#> 
#> 
#> $out
#> AnvlArray
#>  1
#>  3
#>  6
#> [ CPUf32{3} ] 
#>