Package 'anvl'

Description

Code transformation framework for R.

Third-Party Licenses

The anvl package itself is MIT-licensed. The CUDA backend dynamically loads NVIDIA software which is not bundled with anvl, but downloaded from NVIDIA's official redistributable channels by the CUDA toolkit R package (e.g. cuda12.8) at install time. Its use is governed by the NVIDIA CUDA Toolkit EULA, with the exception of cuDNN, which is covered by the NVIDIA cuDNN SLA, and NCCL, which is covered by its own license. By installing or using the CUDA backend you accept those terms.

Author(s)

Maintainer: Sebastian Fischer [email protected] (ORCID)

Authors:

• Daniel Falbel [email protected] (ORCID)

• Tomasz Kalinowski [email protected]

• Nikolai German [email protected] (ORCID)

See Also

Useful links:

• https://r-xla.github.io/anvl/

• https://github.com/r-xla/anvl

• Report bugs at https://github.com/r-xla/anvl/issues

Description

Get the current graph being built (via local_descriptor).

Usage

.current_descriptor(silent = FALSE)

Arguments

Value

A GraphDescriptor object.

Description

Extracts a subset from an array. You can also use the [ operator. Supports R-style indexing including scalar indices (which drop dimensions), ranges (a:b), and array(c(...)) for selecting multiple elements along a dimension.

Usage

## S3 method for class 'AnvlArray'
x[...]

nv_subset(operand, ...)

Arguments

Value

arrayish

See Also

nv_subset_assign() for updating subsets, vignette("subsetting") for a comprehensive guide.

Examples

x <- nv_matrix(1:12, nrow = 3)
x
# Select row 2
x[2, ]

# Select rows 1 to 2, all columns
x[1:2, ]

Description

Updates elements of an array at specified positions, returning a new array. You can also use the ⁠[<-⁠ operator.

Usage

## S3 replacement method for class 'AnvlArray'
x[...] <- value

nv_subset_assign(operand, ..., value)

Arguments

Value

arrayish
A new array with the same shape as operand and the subset replaced.

See Also

nv_subset(), vignette("subsetting") for a comprehensive guide.

Examples

x <- nv_matrix(1:12, nrow = 3)
# Set row 1 to zeros
x[1, ] <- nv_scalar(0L)
x

Description

Calls the extractor after converting the input to an AbstractArray.

Usage

shape_abstract(x)

ndims_abstract(x)

dtype_abstract(x)

ambiguous_abstract(x)

Arguments

Description

Returns whether the array's dtype is ambiguous.

Usage

ambiguous(x, ...)

Arguments

Value

logical(1) - TRUE if the dtype is ambiguous, FALSE otherwise

Description

The main array object. Its type is determined by a data type and a shape.

Usage

nv_array(
  data,
  dtype = NULL,
  device = NULL,
  shape = NULL,
  ambiguous = NULL,
  backend = NULL,
  byrow = FALSE
)

nv_scalar(data, dtype = NULL, device = NULL, ambiguous = NULL, backend = NULL)

nv_matrix(
  data,
  nrow = NULL,
  ncol = NULL,
  dtype = NULL,
  device = NULL,
  ambiguous = NULL,
  backend = NULL,
  byrow = FALSE
)

nv_empty(dtype, shape, device = NULL, ambiguous = FALSE)

nv_array_like(
  like,
  data,
  dtype = NULL,
  device = NULL,
  shape = NULL,
  ambiguous = NULL,
  backend = NULL
)

nv_scalar_like(
  like,
  data,
  dtype = NULL,
  device = NULL,
  ambiguous = NULL,
  backend = NULL
)

nv_empty_like(
  like,
  dtype = NULL,
  shape = NULL,
  device = NULL,
  ambiguous = NULL
)

Arguments

Value

(AnvlArray)

Extractors

The following generic functions can be used to extract information from an AnvlArray:

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

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

• ndims(): Get the number of dimensions.

• device(): Get the device of the array.

• platform(): Get the platform (e.g. "cpu", "cuda").

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

Serialization

Arrays can be serialized to and from the safetensors format:

• nv_save() / nv_read(): Save/load arrays to/from a file.

• nv_serialize() / nv_unserialize(): Serialize/deserialize arrays to/from raw vectors.

See Also

nv_fill, nv_iota, nv_seq, as_array, nv_serialize

Examples

# A 1-d array (vector) with shape (4). Default type for integers is `i32`
nv_array(1:4)

# Specify a dtype
nv_array(c(1.5, 2.5, 3.5), dtype = "f64")

# A 2x3 matrix
nv_array(1:6, shape = c(2L, 3L))

# A 2x3 matrix filled by row, like `matrix(1:6, 2, 3, byrow = TRUE)`.
nv_array(1:6, shape = c(2L, 3L), byrow = TRUE)

# A scalar array.
nv_scalar(3.14)

# A 0x3 array
nv_empty("f32", shape = c(0L, 3L))

# --- Extractors ---
x <- nv_array(1:6, shape = c(2L, 3L))
dtype(x)
shape(x)
ndims(x)
device(x)
platform(x)
ambiguous(x)

# --- Transforming arrays with jit ---
add_one <- jit(function(x) x + 1)
add_one(nv_array(1:4))

# --- Eager mode (calling operations directly) ---
nv_add(nv_array(1:3), nv_array(4:6))

Description

Constructs the quickr backend, which stores array data as plain R arrays and compiles jitted functions to R code via the quickr package.

Usage

AnvlBackendQuickr()

Details

To use it, the "quickr" package needs to be installed.

Registered automatically under the name "quickr" when the package is loaded; call local_backend("quickr") or with_backend("quickr", ...) to use it. Requires the quickr package to be installed.

Value

An AnvlBackend object with subclass "AnvlBackendQuickr".

Data representation

An AnvlArray with backend = "quickr" is, under the hood, a plain R vector or array (numeric, integer, or logical) stored in the ⁠$data⁠ field. as_array() returns the underlying vector/array directly without copying, and nv_array() simply wraps an R vector/array. As a consequence, there is no separate notion of a device: data always lives in R's memory and computation always runs on the CPU.

Status

This backend is experimental and has a number of limitations:

• Compilation (tracing + quickr lowering) is somewhat slow, so it is best suited to long-running or repeatedly-called functions where the one-time compilation cost is amortized.

• Only a subset of the primitives that the XLA backend supports are currently lowered to quickr code. See vignette("primitives") for an overview.

• Only the data types f64, i32, and bool are supported.

• Only CPU execution is supported.

Quickr JIT arguments

• unwrap (logical(1), default FALSE): if TRUE, the compiled function returns plain R arrays instead of AnvlArrays. Useful when the jitted function's output is consumed by non-anvl R code and the extra wrapping would only get stripped again.

See Also

AnvlBackend(), AnvlBackendXla(), local_backend(), jit().

Description

Constructs the XLA backend, which stores array data in PJRT buffers (via pjrt::pjrt_buffer()) and compiles jitted functions to XLA executables via stablehlo() and pjrt::pjrt_compile(). This is the default backend.

Usage

AnvlBackendXla()

Value

An AnvlBackend object with subclass "AnvlBackendXla".

Data representation

An AnvlArray with backend = "xla" wraps a pjrt::pjrt_buffer() stored in the ⁠$data⁠ field. The buffer owns the memory holding the tensor values and may live on any device supported by PJRT (CPU, CUDA, Metal, ...). Calling as_array() transfers the buffer contents back to an R array; calling nv_array() on an R object uploads it to the requested device.

Each AnvlArray therefore has an associated device, queryable via device(). A device is a pjrt::as_pjrt_device() object (e.g. the platform "cpu" or "cuda", optionally with an index such as "cuda:1"). When device is NULL in nv_array() or the jit() wrapper, the device defaults to the PJRT_PLATFORM environment variable (falling back to "cpu"), or is inferred from the existing inputs of a jitted call. Operations require all inputs to live on the same device.

XLA JIT arguments

• donate (character(), default character()): names of arguments whose underlying buffers may be donated to (i.e., reused/consumed by) the compiled XLA executable. Donated buffers must not be used again by the caller after the call; this can reduce memory usage and copies for large inputs. Must not overlap with static.

See Also

AnvlBackend(), AnvlBackendQuickr(), local_backend(), jit().

Description

Virtual S3 base class for GraphBox.

See Also

GraphBox

Description

Computational graph consisting exclusively of primitive calls. This is a mutable class.

Usage

AnvlGraph(
  calls = list(),
  in_tree = NULL,
  out_tree = NULL,
  inputs = list(),
  outputs = list(),
  constants = list(),
  is_static_flat = NULL,
  static_args_flat = NULL
)

Arguments

Value

(AnvlGraph)

Description

Primitive interpretation rule. Note that [[ and ⁠[[<-⁠ access the interpretation rules. To access other fields, use $ and ⁠$<-⁠.

A primitive is considered higher-order if it contains subgraphs.

Usage

AnvlPrimitive(name, subgraphs = character())

Arguments

Value

(AnvlPrimitive)

Description

Create an R array without having to wrap data in c()

Usage

arr(..., shape = NULL)

Arguments

Examples

arr(1, 2, 3)
arr(1, 2, 3, 4, shape = c(2, 2))

Description

A arrayish value is any object that can be input to a primitive such as prim_add.

During runtime of a JIT-compiled function, these are AnvlArray objects.

The following types are arrayish (during tracing):

• AnvlArray: a concrete array holding data on a device.

• GraphBox: a boxed abstract array representing a value in a graph.

• Length-1 vectors: numeric(1) and logical(1)

• R arrays of types: numeric and logical.

Use is_arrayish() to check whether a value is arrayish.

Usage

is_arrayish(x, convert_ok = TRUE)

Arguments

Value

logical(1)

See Also

AnvlArray, GraphBox

Examples

# AnvlArrays are arrayish
is_arrayish(nv_array(1:4))

# Scalar R literals are arrayish by default
is_arrayish(1.5)
# R arrays are arrayish by default
is_arrayish(array(1.5))

# R arrays
is_arrayish(array(1:4), convert_ok = TRUE)
is_arrayish(array(1:4), convert_ok = FALSE)

# Length 1 vectors
is_arrayish(1.5, convert_ok = FALSE)
is_arrayish(1.5, convert_ok = TRUE)

Description

Use this to canonicalize inputs at the start of a function so it works both with eager executing and in combination with jit(). Use as_anvl_array() for a single input and as_anvl_arrays() for multiple inputs. The latter will also ensure all arrays are from the same backend and live on the same device.

Usage

as_anvl_array(x, device = NULL)

as_anvl_arrays(...)

Arguments

Details

During tracing, boxes are returned as is. During eager mode, R literals and arrays are converted to AnvlArrays on the specified device. For AnvlArray inputs, we check that they live on provided device (if specified).

Value

(AnvlArray for as_anvl_array(), list of AnvlArrays for as_anvl_arrays()).

Examples

as_anvl_array(1L)
as_anvl_arrays(nv_array(1:3), 1L)

Description

Transfers array data to R and returns it as an R array. Only in the case of scalars is the result a vector of length 1, as R arrays cannot have 0 dimensions.

Usage

as_array(x, ...)

Arguments

Details

This is implemented via the generic tengen::as_array().

Value

An R array or vector of length 1.

Examples

x <- nv_array(1:4, dtype = "f32")
as_array(x)
y <- nv_scalar(1L)
# R arrays can't have 0 dimensions:
as_array(y)

Description

Coerces a value to a DataType. Accepts data type strings (e.g. "f32", "i64", "bool") or existing DataType objects (they are returned unchanged).

Usage

as_dtype(x)

Arguments

Details

This is implemented via the generic tengen::as_dtype().

Value

A DataType object.

See Also

is_dtype(), tengen::as_dtype(), tengen::DataType

Examples

as_dtype("f32")
as_dtype("i32")

Description

Returns the underlying bytes of an array as a raw vector.

Usage

as_raw(x, ...)

Arguments

Details

This is implemented via the generic tengen::as_raw().

Value

A raw vector.

Examples

x <- nv_array(1:4, shape = c(2, 2), dtype = "f32")
as_raw(x, row_major = TRUE)
as_raw(x, row_major = FALSE)

Description

Convert an AnvlArray to a bare R vector. The array's shape is discarded; the result is always a flat vector. Each method requires a compatible dtype:

• as.double() / as.numeric(): float or (signed/unsigned) integer dtypes.

• as.integer(): signed or unsigned integer dtypes.

• as.logical(): bool.

• as.vector(): any dtype; the R type is chosen by the dtype, or forced via mode (e.g. "integer", "double", "logical", "list").

Use as_array() to obtain an R array that preserves the shape, or nv_convert() to change the dtype of an AnvlArray before coercing.

Usage

## S3 method for class 'AnvlArray'
as.double(x, ...)

## S3 method for class 'AnvlArray'
as.integer(x, ...)

## S3 method for class 'AnvlArray'
as.logical(x, ...)

## S3 method for class 'AnvlArray'
as.vector(x, mode = "any")

Arguments

Value

An R vector of the corresponding type (double, integer, or logical).

Examples

x <- nv_array(c(1.5, 2.5, 3.5, 4.5), shape = c(2L, 2L))
as.numeric(x)
as.integer(nv_array(1:6, shape = c(2L, 3L)))
as.logical(nv_array(c(TRUE, FALSE), dtype = "bool"))
as.vector(x)

Description

Convert an AbstractArray to a ValueType.

Usage

at2vt(x)

Arguments

Value

(ValueType)

Description

Block until the array's underlying computation has finished, and return the object invisibly. Useful for benchmarking, where the dispatch of an asynchronous operation should not be confused with its execution.

Usage

await(x, ...)

Arguments

Details

Implemented via the generic pjrt::await(). For backends without asynchronous execution (e.g. "quickr"), this is a no-op.

Value

x, invisibly.

See Also

pjrt::await(), map_tree() (to await a tree of outputs)

Examples

x <- nv_array(1:4, dtype = "f32")
await(x)

# Await all leaves of a (possibly nested) list of arrays.
map_tree(list(x, list(y = x)), await)

Description

Get Backend of an Array

Usage

backend(x, ...)

Arguments

Value

character(1) - the backend name

Description

Captures the nesting structure of an object as a tree of Nodes. Each leaf in the input becomes a LeafNode with an integer index corresponding to its position in the flat list produced by flatten(). Lists become ListNodes that record child nodes and names. The resulting tree can be passed to unflatten() to reconstruct the original structure from a flat list.

Usage

build_tree(x, counter = NULL)

Arguments

Value

A Node (LeafNode for scalars, ListNode for lists).

See Also

flatten(), unflatten(), tree_size(), reindex_tree()

Examples

x <- list(a = 1, b = list(c = 2, d = 3))
tree <- build_tree(x)
tree_size(tree)

flat <- flatten(x)
unflatten(tree, flat)

Description

Computes the common dtype for a set of abstract types, respecting whether a type is ambiguous or not. A type is ambiguous if it comes from a literal (like 1 or 1.0) or was promoted to an ambiguous type. Promoting to an ambiguous type can happen in scenarios like x + 1.2, where x is a bool or an int.

Usage

common_dtype(
  lhs_dtype,
  rhs_dtype,
  lhs_ambiguous = FALSE,
  rhs_ambiguous = FALSE
)

Arguments

Value

(⁠list(dtype = [⁠tengen::DataType⁠], ambiguous = ⁠logical(1)')

Description

An AbstractArray that also holds a reference to the actual array data. Usually represents a closed-over constant in a program. Inherits from AbstractArray.

Usage

ConcreteArray(data)

Arguments

Lowering

When lowering to XLA, these become inputs to the executable instead of embedding them into programs as constants. This is to avoid increasing compilation time and bloating the size of the executable.

Examples

y <- nv_array(c(0.5, 0.6))
x <- ConcreteArray(y)
x
ambiguous(x)
shape(x)
ndims(x)
dtype(x)

# How it appears during tracing
graph <- trace_fn(function() y, list())
graph
graph$outputs[[1]]$aval

Description

Returns the target platform currently set by an enclosing stablehlo() call (e.g. "cpu", "cuda"). Platform-aware lowering rules call this to branch on the target — e.g. SVD switches to a layout-flip variant when targeting CUDA with m < n because cuSOLVER's gesvd requires m >= n. Returns NULL outside of a lowering call.

local_platform() sets the current platform for the duration of the calling scope, restoring the previous value via withr::defer() when the scope exits. Useful in tests and for manually exercising platform-aware lowering rules outside of a stablehlo() call.

Usage

current_platform()

local_platform(platform, envir = parent.frame())

Arguments

Value

current_platform() returns NULL or character(1). local_platform() invisibly returns the previous platform.

See Also

stablehlo()

Description

Returns the current default backend from getOption("anvl.default_backend", "xla").

Usage

default_backend()

Value

character(1) — the backend name (e.g. "xla", "quickr").

See Also

local_backend()

Description

Returns a device object for the default backend and platform. For the "xla" backend, the platform is determined by the PJRT_PLATFORM environment variable (defaulting to "cpu"). Other backends (e.g. "quickr") only support CPU. The backend defaults to default_backend().

Usage

default_device(backend = NULL)

Arguments

Value

A backend-specific device object.

See Also

nv_device(), default_backend()

Description

Returns the device on which an array is allocated.

Usage

device(x, ...)

Arguments

Details

This is implemented via the generic tengen::device().

Value

PJRTDevice

Examples

x <- nv_array(1:4, dtype = "f32")
device(x)

Description

Pass the result to jit()'s device argument to indicate that the device should be read from a formal argument of the function being compiled. At call time, the value of that argument is used to derive the backend via backend() dispatch and is forwarded to the backend-specific JIT as the compilation device.

This is intended for functions that have no dynamic array inputs from which the backend could otherwise be detected (e.g. array constructors like prim_fill() or prim_iota()).

Usage

device_arg(argname)

Arguments

Value

(AnvlDeviceArg)
An object recognized by jit().

See Also

jit(), backend()

Examples

f <- function(x) nv_scalar(1, device = x)
g <- jit(f, backend = "auto", device = device_arg("x"))
g(nv_device("cpu", "xla"))

Description

Returns the data type of an array (e.g. f32, i64).

Usage

dtype(x, ...)

Arguments

Details

This is implemented via the generic tengen::dtype().

Value

A DataType.

See Also

tengen::dtype()

Examples

x <- nv_array(1:4, dtype = "f32")
dtype(x)

Description

Compare two abstract arrays for type equality.

Usage

eq_type(e1, e2, ambiguity)

neq_type(e1, e2, ambiguity)

Arguments

Value

logical(1) - TRUE if the arrays are equal, FALSE otherwise.

Examples

a <- nv_aval("f32", c(2L, 3L))
b <- nv_aval("f32", c(2L, 3L))

# Same dtype and shape
eq_type(a, b, ambiguity = FALSE)

# Different dtype
eq_type(a, nv_aval("i32", c(2L, 3L)), ambiguity = FALSE)

# Different shape
eq_type(a, nv_aval("f32", c(3L, 2L)), ambiguity = FALSE)

# ambiguity parameter controls whether ambiguous field is compared
c <- nv_aval("f32", c(2L, 3L), ambiguous = TRUE)
eq_type(a, c, ambiguity = FALSE)
eq_type(a, c, ambiguity = TRUE)

# neq_type is the negation of eq_type
neq_type(a, b, ambiguity = FALSE)

Description

Subsets a ListNode to keep only the children whose names match names, then reindexes the leaf nodes so they map to contiguous positions in a flat list. If all names are kept the original tree is returned unchanged.

Usage

filter_list_node(tree, names)

Arguments

Value

A ListNode containing only the selected children with reindexed leaves.

See Also

build_tree(), reindex_tree(), unflatten()

Examples

x <- list(a = 1, b = 2, c = 3)
tree <- build_tree(x)
sub <- filter_list_node(tree, c("a", "c"))
tree_size(sub)

unflatten(sub, x[c("a", "c")])

Description

Recursively flattens a nested list into a single flat list containing only the leaf values, preserving left-to-right order.

Currently only lists are flattened and all other objects are treated as leaves.

Use build_tree() to capture the nesting structure so it can be restored with unflatten().

Usage

flatten(x)

Arguments

Value

list() containing the flattened values.

See Also

build_tree(), unflatten(), tree_size()

Examples

x <- list(a = 1, b = list(c = 2, d = 3))
flatten(x)

flatten(list(1:3, "hello"))

Description

Returns a new function that computes the gradient of f via reverse-mode automatic differentiation. f must return a single float scalar. The returned function has the same signature as f and returns the gradients in the same structure as the inputs (or the subset selected by wrt).

Usage

gradient(f, wrt = NULL)

Arguments

Value

function

See Also

value_and_gradient() to get both the output and gradients, transform_gradient() for the low-level graph transformation.

Examples

f <- function(x, y) sum(x * y)
g <- jit(gradient(f))
g(nv_array(c(1, 2), dtype = "f32"), nv_array(c(3, 4), dtype = "f32"))

# Differentiate with respect to a single argument
g_x <- jit(gradient(f, wrt = "x"))
g_x(nv_array(c(1, 2), dtype = "f32"), nv_array(c(3, 4), dtype = "f32"))

# Static (non-array) arguments are passed through but cannot be in wrt
f2 <- function(x, power) sum(x^power)
g2 <- jit(gradient(f2, wrt = "x"), static = "power")
g2(nv_array(c(1, 2, 3), dtype = "f32"), power = 2L)

Description

Add a primitive call to a graph descriptor. Inside a primitive body created with new_primitive(), pass the lexically-bound self as the primitive argument.

Usage

graph_desc_add(primitive, args, params = list(), infer_fn, desc = NULL)

Arguments

Value

(list of GraphBox)

Description

Lowers a supported subset of AnvlGraph objects to a plain R function (no compilation) suitable for quickr::quick(). The returned function expects plain R scalars/vectors/arrays and returns plain R values/arrays.

Usage

graph_to_quickr_r_function(graph)

Arguments

Details

Most users will prefer jit() with backend = "quickr". This function is the lower-level graph API.

Value

(function)

See Also

jit() with options(anvl.backend = "quickr") for tracing and compiling a regular R function in one step.

Description

An AnvlBox subclass that wraps a GraphNode during graph construction (tracing). When a function is traced via trace_fn(), each intermediate array value is represented as a GraphBox. It also contains an associated GraphDescriptor in which the node "lives".

Usage

GraphBox(gnode, desc)

Arguments

Value

(GraphBox)

Extractors

• dtype()

• shape()

• ndims()

• ambiguous()

See Also

AnvlBox, trace_fn(), jit()

Description

Descriptor of an AnvlGraph. This is a mutable class.

Usage

GraphDescriptor(
  calls = list(),
  tensor_to_gval = NULL,
  gval_to_box = NULL,
  constants = list(),
  in_tree = NULL,
  out_tree = NULL,
  inputs = list(),
  outputs = list(),
  is_static_flat = NULL,
  static_args_flat = NULL,
  devices = character()
)

Arguments

Value

(GraphDescriptor)

Description

Literal in an AnvlGraph. This is a mutable class.

Usage

GraphLiteral(aval)

Arguments

Value

(GraphLiteral)

Description

Virtual base class for nodes in an AnvlGraph. Is either a GraphValue or a GraphLiteral. Cannot be instantiated directly - use GraphValue() or GraphLiteral() instead.

Description

Value in an AnvlGraph. This is a mutable class.

Usage

GraphValue(aval)

Arguments

Value

(GraphValue)

Description

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, dimension, start = 1L, ambiguous = FALSE)

Arguments

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 stablehlo::hlo_iota(), optionally shifting the starting value via stablehlo::hlo_add().

Examples

x <- IotaArray(shape = 4L, dtype = "i32", dimension = 1L)
x
ambiguous(x)
shape(x)
ndims(x)
dtype(x)
# How it appears during tracing:
graph <- trace_fn(function() nv_iota(dim = 1L, dtype = "i32", shape = 4L), list())
graph
graph$outputs[[1]]$aval

Description

Test whether an object is a device

Usage

is_device(x)

Arguments

Value

logical(1)

Description

Tests whether x is a DataType object.

Usage

is_dtype(x)

Arguments

Value

TRUE or FALSE.

See Also

as_dtype(), tengen::is_dtype()

Examples

is_dtype("f32")
is_dtype(as_dtype("f32"))

Description

Wraps a function so that it is traced and compiled on first call. Subsequent calls with the same input structure, shapes, and dtypes hit an LRU cache and skip recompilation. Unlike xla(), the compiled executable is not created eagerly but lazily on the first invocation.

Usage

jit(
  f,
  static = character(),
  cache_size = 100L,
  backend = NULL,
  device = NULL,
  ...
)

Arguments

Value

A JitFunction (a function with the same formals as f). The returned wrapper expects AnvlArray inputs and returns AnvlArray values.

Device and Backend selection

There are various ways to specify which device and which backend to use.

Concrete backend: In the case where we fix a concrete backend (backend is not "auto"), the device can be inferred or set explicitly. Setting the device explicitly allows you to enforce that the function always uses the specified device, e.g. "cuda:0". If the device argument is set, all encountered arrays are copied to it.

If the device is not specified (NULL; default) the device will be inferred from the input arrays and the constants within the program. If conflicting devices are found, an error is thrown. If no array with a device is found, we fall back to the default device.

Auto backend: When setting backend = "auto", the backend will be inferred from the array inputs and otherwise fall back to the default backend. If you want to jit() a function without array inputs but make it work with different devices, set device = device_arg("<argname>") where ⁠<argname>⁠ is the name of the argument specifying the device. Note that this is only necessary with the "auto" backend. When using a concrete backend, you can just specify the device via a static argument.

XLA JIT arguments

• donate (character(), default character()): names of arguments whose underlying buffers may be donated to (i.e., reused/consumed by) the compiled XLA executable. Donated buffers must not be used again by the caller after the call; this can reduce memory usage and copies for large inputs. Must not overlap with static.

Quickr JIT arguments

• unwrap (logical(1), default FALSE): if TRUE, the compiled function returns plain R arrays instead of AnvlArrays. Useful when the jitted function's output is consumed by non-anvl R code and the extra wrapping would only get stripped again.

See Also

xla() for ahead-of-time compilation, jit_eval() for evaluating an expression once.

Examples

f <- jit(function(x, y) x + y)
f(nv_array(1), nv_array(2))

# Static arguments enable data-dependent control flow
g <- jit(function(x, flag) {
  if (flag) x + 1 else x * 2
}, static = "flag")
g(nv_array(3), TRUE)
g(nv_array(3), FALSE)


with_backend("quickr", {
  h <- jit(function(x, y) x + y)
  h(nv_array(1), nv_array(2))
})

Description

Convenience wrapper that JIT-compiles and immediately evaluates a single expression. Equivalent to wrapping expr in an anonymous function, calling jit() on it, and invoking the result. Useful if you want to evaluate an expression once.

Usage

jit_eval(expr, ...)

Arguments

Value

(any)
Result of the compiled and evaluated expression.

Examples

x <- nv_array(c(1, 2, 3), dtype = "f32")
jit_eval(x + x)

Description

An AbstractArray where all elements have the same constant value. This either arises when using literals in traced code (e.g. x + 1) or when using nv_fill() to create a constant.

Usage

LiteralArray(data, shape, dtype = default_dtype(data), ambiguous)

Arguments

Type Ambiguity

When arising from R literals, the resulting LiteralArray is ambiguous because no type information was available. See the vignette("type-promotion") for more details.

Lowering

LiteralArrays become constants inlined into the stableHLO program. I.e., they lower to stablehlo::hlo_tensor().

Examples

x <- LiteralArray(1L, shape = integer(), ambiguous = TRUE)
x
ambiguous(x)
shape(x)
ndims(x)
dtype(x)
# How it appears during tracing:
# 1. via R literals
graph <- trace_fn(function() 1, list())
graph
graph$outputs[[1]]$aval
# 2. via nv_fill()
graph <- trace_fn(function() nv_fill(2L, shape = c(2, 2)), list())
graph
graph$outputs[[1]]$aval

Description

Sets the anvl.default_backend option for the duration of the calling scope. This affects nv_array(), nv_scalar(), and jit().

Usage

local_backend(backend, envir = parent.frame())

Arguments

Value

The previous value of the option (invisibly).

Description

Creates a new GraphDescriptor which is afterwards accessible via .current_descriptor(). The graph is automatically removed when exiting the current scope. After the graph is either cleaned up automatically (by exiting the scope) or finalized, the previously built graph is restored, i.e., accessible via .current_descriptor().

Usage

local_descriptor(..., envir = parent.frame())

Arguments

Value

A GraphDescriptor object.

Description

Apply a function to each leaf of a (possibly nested) list, preserving the tree structure. Equivalent to flattening .x with flatten(), applying .f to each leaf, and reassembling with unflatten().

Usage

map_tree(.x, .f, ...)

Arguments

Value

An object with the same nesting structure as .x, where each leaf is the result of .f(leaf, ...).

See Also

flatten(), build_tree(), unflatten(), await()

Examples

map_tree(list(a = 1, b = list(c = 2, d = 3)), \(x) x + 1)

Description

Returns the number of dimensions (sometimes also refered to as rank) of an array. Equivalent to length(shape(x)).

Usage

ndims(x)

Arguments

Value

integer(1)

See Also

tengen::ndims()

Examples

x <- nv_array(1:4, dtype = "f32")
ndims(x)

Description

Builds an AnvlPrimitive metadata object, wraps fn with jit(), attaches the metadata via attr(., "primitive"), prepends class "JitPrimitive", and (by default) registers the result under name in the primitive registry.

The backend is always "auto" and cannot be configured.

Usage

new_primitive(
  name,
  fn,
  subgraphs = character(),
  static = character(),
  device = NULL,
  register = TRUE
)

Arguments

Value

A callable of class c("JitPrimitive", "JitFunction").

Description

Element-wise absolute value. You can also use abs().

Usage

nv_abs(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_abs() for the underlying primitive.

Examples

x <- nv_array(c(-1, 2, -3))
abs(x)

Description

Element-wise inverse cosine. You can also use acos().

Usage

nv_acos(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_acos() for the underlying primitive.

Examples

x <- nv_array(c(-1, 0, 1))
acos(x)

Description

Element-wise inverse hyperbolic cosine. You can also use acosh().

Usage

nv_acosh(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_acosh() for the underlying primitive.

Examples

x <- nv_array(c(1, 2, 10))
acosh(x)

Description

Adds two arrays element-wise. You can also use the + operator.

Usage

nv_add(lhs, rhs)

Arguments

Value

arrayish
Has the same shape and the promoted common data type of the inputs.

See Also

prim_add() for the underlying primitive.

Examples

x <- nv_array(c(1, 2, 3))
y <- nv_array(c(4, 5, 6))
x + y

Description

Element-wise logical AND. You can also use the & operator.

Usage

nv_and(lhs, rhs)

Arguments

Value

arrayish
Has the same shape and the promoted common data type of the inputs.

See Also

prim_and() for the underlying primitive.

Examples

x <- nv_array(c(TRUE, FALSE, TRUE))
y <- nv_array(c(TRUE, TRUE, FALSE))
x & y

Description

Returns the index of the maximum value along a dimension. Ties are broken by returning the smallest index.

Usage

nv_argmax(operand, dim = NULL, drop = TRUE, nan_rm = FALSE)

Arguments

Value

arrayish of dtype i32
Same shape as operand with dim removed (or set to 1 if drop = FALSE).

NaN handling

With nan_rm = FALSE (default), if any entry along the reduced axis is NaN, the returned index points at the first such NaN. With nan_rm = TRUE, NaN entries are skipped.

See Also

nv_argmin(), nv_reduce_max().

Examples

nv_argmax(nv_array(c(3, 1, 4, 1, 5, 9, 2, 6)))
nv_argmax(nv_matrix(c(3, 1, 5, 2, 4, 0), nrow = 2, byrow = TRUE),
  dim = 2L
)
nv_argmax(nv_array(c(1, NaN, 3)))
nv_argmax(nv_array(c(1, NaN, 3)), nan_rm = TRUE)

Description

Returns the index of the minimum value along a dimension. Ties are broken by returning the smallest index.

Usage

nv_argmin(operand, dim = NULL, drop = TRUE, nan_rm = FALSE)

Arguments

Value

arrayish of dtype i32
Same shape as operand with dim removed (or set to 1 if drop = FALSE).

NaN handling

With nan_rm = FALSE (default), if any entry along the reduced axis is NaN, the returned index points at the first such NaN. With nan_rm = TRUE, NaN entries are skipped.

See Also

nv_argmax(), nv_reduce_min().

Examples

nv_argmin(nv_array(c(3, 1, 4, 1, 5, 9, 2, 6)))
nv_argmin(nv_array(c(2, NaN, 1, 3)))
nv_argmin(nv_array(c(2, NaN, 1, 3)), nan_rm = TRUE)

Description

Returns the indices that would sort the array along a dimension.

Usage

nv_argsort(operand, dim = NULL, decreasing = FALSE, stable = FALSE)

Arguments

Value

arrayish of dtype i32
Same shape as operand. For a size-0 axis, the output is an empty i32 array of the same shape (a valid empty permutation). as_array(operand)[as_array(nv_argsort(operand))] reproduces the sorted array (for 1-D inputs).

NaN handling

NaN values sort to the end (ascending) or beginning (descending), regardless of sign. +0 and -0 compare equal.

See Also

nv_sort(), prim_sort().

Examples

x <- nv_array(c(3, 1, 4, 1, 5))
nv_argsort(x)

Description

Element-wise inverse sine. You can also use asin().

Usage

nv_asin(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_asin() for the underlying primitive.

Examples

x <- nv_array(c(-1, 0, 1))
asin(x)

Description

Element-wise inverse hyperbolic sine. You can also use asinh().

Usage

nv_asinh(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_asinh() for the underlying primitive.

Examples

x <- nv_array(c(-1, 0, 1))
asinh(x)

Description

Element-wise inverse tangent. You can also use atan().

Usage

nv_atan(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_atan() for the underlying primitive.

Examples

x <- nv_array(c(-1, 0, 1))
atan(x)

Description

Element-wise two-argument arctangent, i.e. the angle (in radians) between the positive x-axis and the point ⁠(rhs, lhs)⁠.

Usage

nv_atan2(lhs, rhs)

Arguments

Value

arrayish
Has the same shape and the promoted common data type of the inputs.

See Also

prim_atan2() for the underlying primitive.

Examples

y <- nv_array(c(1, 0, -1))
x <- nv_array(c(0, 1, 0))
nv_atan2(y, x)

Description

Element-wise inverse hyperbolic tangent. You can also use atanh().

Usage

nv_atanh(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_atanh() for the underlying primitive.

Examples

x <- nv_array(c(-0.5, 0, 0.5))
atanh(x)

Description

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

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

• closed-over constants: ConcreteArray

• scalar arrays arising from R literals: LiteralArray

• sequence patterns: IotaArray

To convert a arrayish value to an abstract array, use to_abstract().

Usage

nv_aval(dtype, shape, ambiguous = FALSE)

AbstractArray(dtype, shape, ambiguous = FALSE)

Arguments

Extractors

The following extractors are available on AbstractArray objects:

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

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

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

• ndims(): Get the number of dimensions.

See Also

LiteralArray, ConcreteArray, IotaArray, GraphValue, to_abstract(), GraphBox

Examples

# -- Creating abstract arrays --
a <- AbstractArray("f32", c(2L, 3L))
a
dtype(a)
shape(a)
ambiguous(a)

# Shorthand
nv_aval("f32", c(2L, 3L))

# How AbstractArrays appear in an AnvlGraph
graph <- trace_fn(function(x) x + 1, list(x = nv_aval("i32", 4L)))
graph
graph$inputs[[1]]$aval

Description

Combine arrays along the row (nv_rbind) or column (nv_cbind) dimension. Arguments are first promoted to a common data type (see nv_promote_to_common()).

Each input is then handled according to its rank:

• 0-D: broadcast to match the non-stacked dimensions of the other inputs.

• 1-D: treated as a single row/column.

• Other: used as-is.

Usage

nv_rbind(...)

nv_cbind(...)

## S3 method for class 'AnvlArray'
rbind(..., deparse.level = 1)

## S3 method for class 'AnvlArray'
cbind(..., deparse.level = 1)

Arguments

Value

arrayish

Differences from base R

base::rbind() and base::cbind() applied to an array() of rank > 2 flatten the trailing dimensions into the column axis (so a c(2, 3, 4) array becomes a ⁠2 x 12⁠ matrix). nv_rbind and nv_cbind instead preserve all non-stacked dimensions: combining two c(2, 3, 4) arrays with nv_rbind produces a c(4, 3, 4) array, and with nv_cbind a c(2, 6, 4) array.

See Also

nv_concatenate()

Examples

# Vectors as rows / columns
nv_rbind(nv_array(1:3), nv_array(4:6))
nv_cbind(nv_array(1:3), nv_array(4:6))

# Scalar broadcasting
nv_rbind(nv_matrix(1:6, nrow = 2), nv_scalar(0))

# Rank-3 arrays preserve trailing dimensions
a <- nv_array(1:24, shape = c(2, 3, 4))
shape(nv_rbind(a, a)) # c(4, 3, 4)

Description

Reinterprets the bits of an array as a different data type without modifying the underlying data. If the target type is narrower, an extra trailing dimension is added; if wider, the last dimension is consumed.

Usage

nv_bitcast_convert(operand, dtype)

Arguments

Value

arrayish
Has the given dtype.

See Also

prim_bitcast_convert() for the underlying primitive, nv_convert() for value-preserving type conversion.

Examples

x <- nv_array(1L)
prim_bitcast_convert(x, dtype = "i8")

Description

Broadcasts arrays to a common shape using NumPy-style broadcasting rules.

Usage

nv_broadcast_arrays(...)

Arguments

Value

(list() of arrayish)
List of arrays, all with the same shape.

Broadcasting Rules

1. If the arrays have different numbers of dimensions, prepend size-1 dimensions to the shorter shape.

2. For each dimension: if the sizes match, keep them; if one is 1, expand it to the other's size; otherwise raise an error.

See Also

nv_broadcast_scalars(), nv_broadcast_to()

Examples

x <- nv_matrix(1:6, nrow = 2)
y <- nv_array(c(10, 20, 30))
nv_broadcast_arrays(x, y)

Description

Broadcast scalar arrays to match the shape of non-scalar arrays. All non-scalar arrays must have the same shape.

Usage

nv_broadcast_scalars(...)

Arguments

Value

(list() of arrayish)
List of broadcasted arrays.

Examples

x <- nv_array(c(1, 2, 3))
# scalar 1 is broadcast to shape [3]
nv_broadcast_scalars(x, nv_scalar(1))

Description

Broadcasts an array to a target shape using NumPy-style broadcasting rules.

Usage

nv_broadcast_to(operand, shape)

Arguments

Value

arrayish
Has the given shape and the same data type as operand.

See Also

nv_broadcast_arrays(), nv_broadcast_scalars(), prim_broadcast_in_dim() for the underlying primitive.

Examples

x <- nv_array(c(1, 2, 3))
nv_broadcast_to(x, shape = c(2, 3))

Description

Element-wise cube root.

Usage

nv_cbrt(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_cbrt() for the underlying primitive.

Examples

x <- nv_array(c(1, 8, 27))
nv_cbrt(x)

Description

Element-wise ceiling (round toward positive infinity). You can also use ceiling().

Usage

nv_ceiling(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_ceil() for the underlying primitive.

Examples

x <- nv_array(c(1.2, 2.7, -1.5))
ceiling(x)

Description

Computes the Cholesky decomposition of a symmetric positive-definite matrix. Supports batched inputs: dimensions before the last two are batch dimensions.

Usage

nv_chol(operand, lower = FALSE)

## S3 method for class 'AnvlArray'
chol(x, ..., lower = FALSE)

Arguments

Value

arrayish
Triangular matrix with the same shape and data type as the input.

See Also

nv_solve(), prim_chol()

Examples

a <- nv_matrix(c(4, 2, 2, 3), nrow = 2, dtype = "f32")
nv_chol(a)

Description

Element-wise clamp: max(min_val, min(operand, max_val)). Converts min_val and max_val to the data type of operand.

Usage

nv_clamp(min_val, operand, max_val)

Arguments

Details

The underlying stableHLO function already broadcasts scalars, so no need to broadcast manually.

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_clamp() for the underlying primitive.

Examples

x <- nv_array(c(-1, 0.5, 2))
nv_clamp(nv_scalar(0), x, nv_scalar(1))

Description

Concatenates arrays along a dimension. Operands are promoted to a common data type and scalars are broadcast before concatenation.

Usage

nv_concatenate(..., dimension = NULL)

Arguments

Value

arrayish
Has the common data type and a shape matching the inputs in all dimensions except dimension, which is the sum of input sizes.

See Also

prim_concatenate() for the underlying primitive.

Examples

x <- nv_array(c(1, 2, 3))
y <- nv_array(c(4, 5, 6))
nv_concatenate(x, y)

Description

Converts the elements of an array to a different data type. Returns the input unchanged if it already has the target type.

Usage

nv_convert(operand, dtype)

Arguments

Value

arrayish
Has the given dtype and the same shape as operand.

See Also

prim_convert() for the underlying primitive.

Examples

x <- nv_array(c(1L, 2L, 3L))
nv_convert(x, dtype = "f32")

Description

Element-wise cosine. You can also use cos().

Usage

nv_cos(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_cos() for the underlying primitive.

Examples

x <- nv_array(c(0, pi / 2, pi))
cos(x)

Description

Element-wise hyperbolic cosine. You can also use cosh().

Usage

nv_cosh(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_cosh() for the underlying primitive.

Examples

x <- nv_array(c(-1, 0, 1))
cosh(x)

Description

Computes t(lhs) %*% rhs. If rhs is missing, computes t(lhs) %*% lhs.

Usage

nv_crossprod(lhs, rhs = NULL)

## S3 method for class 'AnvlArray'
crossprod(x, y = NULL, ...)

Arguments

Value

arrayish

See Also

nv_tcrossprod(), nv_matmul()

Examples

x <- nv_matrix(1:6, nrow = 3, dtype = "f32")
nv_crossprod(x)

Description

Running maximum, optionally along a single dimension.

Usage

nv_cummax(operand, dim = NULL, with_indices = FALSE, nan_rm = FALSE)

Arguments

Value

arrayish (when with_indices = FALSE) or named list of two arrays (when with_indices = TRUE).

Relation to base R

Both nv_cummax() (with dim = NULL) and base::cummax() flatten a multi-dimensional input to 1-D before accumulating, but the flatten order differs: anvl arrays are row-major (C order), so the flattened sequence iterates the last dim fastest, whereas base R uses column-major (Fortran) order. The two agree on 1-D inputs.

See Also

prim_cummax() for the underlying primitive.

Examples

x <- nv_matrix(c(3, 1, 4, 1, 5, 9), nrow = 2)
nv_cummax(x)
nv_cummax(x, dim = 1L)
nv_cummax(x, dim = 1L, with_indices = TRUE)
nv_cummax(nv_array(c(1, NaN, 3)))                # NaN propagates
nv_cummax(nv_array(c(1, NaN, 3)), nan_rm = TRUE) # NaN skipped

Description

Running minimum, optionally along a single dimension.

Usage

nv_cummin(operand, dim = NULL, with_indices = FALSE, nan_rm = FALSE)

Arguments

Value

arrayish (when with_indices = FALSE) or named list of two arrays (when with_indices = TRUE).

Relation to base R

Both nv_cummin() (with dim = NULL) and base::cummin() flatten a multi-dimensional input to 1-D before accumulating, but the flatten order differs: anvl arrays are row-major (C order), so the flattened sequence iterates the last dim fastest, whereas base R uses column-major (Fortran) order. The two agree on 1-D inputs.

See Also

prim_cummin() for the underlying primitive.

Examples

x <- nv_matrix(c(3, 1, 4, 1, 5, 9), nrow = 2)
nv_cummin(x)
nv_cummin(x, dim = 1L)
nv_cummin(x, dim = 1L, with_indices = TRUE)
nv_cummin(nv_array(c(3, NaN, 1)))                # NaN propagates
nv_cummin(nv_array(c(3, NaN, 1)), nan_rm = TRUE) # NaN skipped

Description

Cumulative product, optionally along a single dimension.

Usage

nv_cumprod(operand, dim = NULL, nan_rm = FALSE)

Arguments

Value

arrayish
Has the same shape and data type as the input.

Relation to base R

Both nv_cumprod() (with dim = NULL) and base::cumprod() flatten a multi-dimensional input to 1-D before accumulating, but the flatten order differs: anvl arrays are row-major (C order), so the flattened sequence iterates the last dim fastest, whereas base R uses column-major (Fortran) order. The two agree on 1-D inputs.

See Also

prim_cumprod() for the underlying primitive.

Examples

x <- nv_matrix(1:6, nrow = 2)
nv_cumprod(x)              # row-major flatten, then accumulate
nv_cumprod(x, dim = 1L)    # accumulate along rows
nv_cumprod(nv_array(c(2, NaN, 3)))                # NaN propagates
nv_cumprod(nv_array(c(2, NaN, 3)), nan_rm = TRUE) # NaN treated as 1

Description

Cumulative sum, optionally along a single dimension.

Usage

nv_cumsum(operand, dim = NULL, nan_rm = FALSE)

Arguments

Value

arrayish
Has the same shape and data type as the input.

Relation to base R

Both nv_cumsum() (with dim = NULL) and base::cumsum() flatten a multi-dimensional input to 1-D before accumulating, but the flatten order differs: anvl arrays are row-major (C order), so the flattened sequence iterates the last dim fastest, whereas base R uses column-major (Fortran) order. The two agree on 1-D inputs.

See Also

prim_cumsum() for the underlying primitive.

Examples

x <- nv_matrix(1:6, nrow = 2)
nv_cumsum(x)              # row-major flatten, then accumulate
nv_cumsum(x, dim = 1L)    # accumulate along rows
nv_cumsum(nv_array(c(1, NaN, 3)))                # NaN propagates
nv_cumsum(nv_array(c(1, NaN, 3)), nan_rm = TRUE) # NaN treated as 0

Description

Computes the determinant of a square matrix via nv_determinant().

Usage

nv_det(operand)

Arguments

Value

Scalar arrayish with the same dtype as operand.

See Also

nv_determinant(), nv_solve(), prim_lu()

Examples

a <- nv_matrix(c(4, 3, 6, 3), nrow = 2, dtype = "f64")
nv_det(a)

Description

Computes the determinant of a square matrix in the modulus / sign decomposition matching base R's base::determinant(). For the plain scalar determinant, use nv_det().

Usage

nv_determinant(operand, logarithm = TRUE)

## S3 method for class 'AnvlArray'
determinant(x, logarithm = TRUE, ...)

Arguments

Details

For computing the determinant, we use:

$P A = L U$

$\det(L) = 1$

$\det(A) = \det(U) / \det(P) = \mathrm{sign}(P^{-1}) \, \prod_i U_{ii} = \mathrm{sign}(P) \, \prod_i U_{ii}$

Matching base R's det_ge_real, the magnitude is computed in log space when logarithm = TRUE ($\sum_i \log|U_{ii}|$) and as a direct product when logarithm = FALSE ($\prod_i |U_{ii}|$).

Value

Named list with elements modulus and sign, both scalar arrayish with the same dtype as operand. The full determinant is sign * exp(modulus) (with logarithm = TRUE) or sign * modulus (with logarithm = FALSE).

See Also

nv_det(), nv_solve(), prim_lu()

Examples

a <- nv_matrix(c(4, 3, 6, 3), nrow = 2, dtype = "f64")
nv_determinant(a)
nv_determinant(a, logarithm = FALSE)

Description

Constructs a backend-specific device object.

A device identifies a compute resources, such as CPU, or a specific GPU. It is relevant for data allocation (e.g. via nv_array()) but also compilation (jit).

Usage

nv_device(x, backend = NULL)

Arguments

Value

A backend-specific device object (e.g. PJRTDevice for "xla", quickr_device for "quickr").

See Also

backend(), AnvlBackend().

Examples

# Create CPU device for xla backend:
nv_device("cpu", "xla")
# Create CPU device for quickr backend:
nv_device("cpu", "quickr")
# Pass through an existing device:
dev <- nv_device("cpu")
identical(nv_device(dev), dev)

Description

Creates a diagonal matrix from a 1-D array.

Usage

nv_diag(operand)

Arguments

Value

arrayish
An ⁠n x n⁠ matrix with x on the diagonal and zeros elsewhere.

Examples

nv_diag(nv_array(c(1, 2, 3)))

Description

Element-wise digamma function (logarithmic derivative of the gamma function). You can also use digamma().

Usage

nv_digamma(operand)

Arguments

Value

arrayish
Has the same shape and data type as the input.

See Also

prim_digamma() for the underlying primitive.

Examples

x <- nv_array(c(0.5, 1, 2, 5))
digamma(x)