ag_capture: Capture an ag_* computation as a reusable device graph

View source: R/ag_capture.R

ag_captureR Documentation

Capture an ag_* computation as a reusable device graph

Description

Wraps fn so that, on the GPU, its ag_* operations are recorded once per input shape and every later call runs the recorded ggml graph: upload the inputs, one graph compute, download the outputs – no R-level op dispatch. On the CPU the wrapper simply calls fn, so the same code runs on both.

Usage

ag_capture(fn, params = NULL)

Arguments

fn

Function of named ag_tensor arguments returning an ag_tensor or a named list of ag_tensors. Forward computations only (no with_grad_tape).

params

Optional list of ag_tensors (typically the model parameters) and ag_buffers read with ag_get_rows whose device memory the recording may reference; a change re-records.

Details

fn must depend on the data only through its arguments. Anything it computes in R from data, or uploads itself (ag_tensor(f(x))), would be frozen into the graph at the first call; such uploads, any download and any eager launch during the recording are errors. Weights must be resident on the device (e.g. trained or used on the GPU) before the first call.

The recording is redone when the input shapes, compute dtype or matmul precision change, when the persistent device pool is reset (device switch), and when a tensor in params no longer points at the device memory it had at recording time.

Value

A function taking the same named arguments as numeric matrices (a vector is one column) and returning a matrix, or a named list of matrices.

See Also

ag_capture_free

Examples


W <- ag_param(matrix(rnorm(6), 2, 3))
f <- ag_capture(function(x) ag_relu(ag_matmul(W, x)), params = list(W))
f(x = matrix(1, 3, 4))


ggmlR documentation built on Oct. 23, 2026, 1:07 a.m.