Skip to contents

Representation of an abstract tensor type. During tracing, it is wrapped in a GraphNode held by a GraphBox. In the lowered AnvilGraph it is also part of GraphNodes representing the values in the program.

The base class represents an unknown value, but child classes exist for:

To convert a tensorish value to an abstract tensor, use to_abstract().

Usage

nv_aten(dtype, shape, ambiguous = FALSE)

AbstractTensor(dtype, shape, ambiguous = FALSE)

Arguments

dtype

(stablehlo::TensorDataType | character(1))
The data type of the tensor.

shape

(stablehlo::Shape | integer())
The shape of the tensor. Can be provided as an integer vector.

ambiguous

(logical(1))
Whether the type is ambiguous. Ambiguous types usually arise from R literals (e.g., 1L, 1.0) and follow special promotion rules. See the vignette("type-promotion") for more details.

Extractors

The following extractors are available on AbstractTensor objects:

  • dtype(): Get the data type of the tensor.

  • shape(): Get the shape (dimensions) of the tensor.

  • ambiguous(): Get whether the dtype is ambiguous.

  • ndims(): Get the number of dimensions.

Examples

# -- Creating abstract tensors --
a <- AbstractTensor("f32", c(2L, 3L))
a
#> AbstractTensor(dtype=f32, shape=2x3) 
dtype(a)
#> <f32>
shape(a)
#> [1] 2 3
ambiguous(a)
#> [1] FALSE

# Shorthand
nv_aten("f32", c(2L, 3L))
#> AbstractTensor(dtype=f32, shape=2x3) 

# How AbstractTensors appear in an AnvilGraph
graph <- trace_fn(function(x) x + 1, list(x = nv_aten("i32", 4L)))
graph
#> <AnvilGraph>
#>   Inputs:
#>     %x1: i32[4]
#>   Body:
#>     %1: f32?[4] = convert [dtype = f32, ambiguous = TRUE] (%x1)
#>     %2: f32?[4] = broadcast_in_dim [shape = 4, broadcast_dimensions = <any>] (1:f32?)
#>     %3: f32?[4] = add(%1, %2)
#>   Outputs:
#>     %3: f32?[4] 
graph$inputs[[1]]$aval
#> AbstractTensor(dtype=i32, shape=4)