Iota Array Class

An AbstractArray representing an integer sequence. Usually created by nv_iota() / nv_seq(), which both call prim_iota() internally. Inherits from AbstractArray.

Usage

IotaArray(shape, dtype, axis, start = 1L)

Arguments

shape

(stablehlo::Shape | integer())
The shape of the array.

dtype

(tengen::DataType)
The data type.

axis

(integer(1))
The axis along which values increase.

start

(integer(1))
The starting value.

Lowering

When lowering to StableHLO, these become iota operations that generate the integer sequence so they do not need to actually hold the data in the executable, similar to ALTREPs in R. It lowers to hlo_iota(), optionally shifting the starting value via hlo_add().

Examples

x <- IotaArray(shape = 4L, dtype = "i32", axis = 1L)
x
#> IotaArray(shape=(4), dtype=i32, axis=1, start=1) 
shape(x)
#> [1] 4
naxes(x)
#> [1] 1
dtype(x)
#> <i32>
# how it appears during tracing:
graph <- trace_fn(function() nv_iota(axis = 1L, dtype = "i32", shape = 4L), list())
graph
#> <AnvlGraph> () {
#>   %1: i32[4] = iota [axis = 1, dtype = "i32", shape = 4, start = 1] ()
#>   return %1
#> }
graph$outputs[[1]]$aval
#> IotaArray(shape=(4), dtype=i32, axis=1, start=1)