ag_argmax: Column-wise argmax

View source: R/ag_select.R

ag_argmaxR Documentation

Column-wise argmax

Description

Index of the largest entry in every column of x, for action selection over a [n_actions, batch] matrix of logits or log-probabilities.

Usage

ag_argmax(x)

Arguments

x

ag_tensor or numeric matrix, [n, batch].

Details

Ties go to the FIRST index on every path (R, CPU kernel, Vulkan shader), matching which.max. NaN is never selected; a column with nothing above -Inf (all NaN / all -Inf) gives 0.

Not differentiable. On the GPU this reads the result back to the host, so it is a synchronisation point: fine where the environment needs the action, but keep it out of the training step.

Value

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

Examples

ag_argmax(matrix(c(1, 3, 3,  5, 2, 5), 3, 2))   # 1 0

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