ag_sample_categorical: Sample one index per column from logits

In ggmlR: 'GGML' Tensor Operations for Machine Learning

Sample one index per column from logits

Description

Categorical sampling by the Gumbel-max trick: argmax(x + G) with G = -log(-log(u)), u ~ U(0, 1), draws index i of column j with probability softmax(x[, j])[i]. x may be logits or log-probabilities (they differ by a per-column constant).

Usage

ag_sample_categorical(x)

Arguments

x	ag_tensor or numeric matrix, [n, batch].

Details

The noise comes from stats::runif, so set.seed reproduces the draw, and the CPU and GPU paths see the same noise. On the GPU only the noise goes up and only the indices come down: x stays on the device, and the add and argmax are one launch. Like ag_argmax it is a synchronisation point and not differentiable; get the log-probability of the drawn action with ag_gather(ag_log_softmax(x), a).

Mask illegal entries as for ag_log_softmax: add a large but FINITE negative number (-1e4). Gumbel noise exceeds 1e4 with probability about exp(-1e4), so a masked entry is never drawn while one column entry is unmasked.

Near-ties of x + G are rounded in f32 on the GPU and in double on the CPU, so the two paths can differ in a rare column even with the same seed.

Value

Integer vector of length ncol(x), 0-based.

Examples

set.seed(1)
ag_sample_categorical(matrix(c(0, 0, -1e4,  5, 0, 0), 3, 2))

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