ag_log_softmax: Column-wise log-softmax

In ggmlR: 'GGML' Tensor Operations for Machine Learning

Column-wise log-softmax

Description

log(softmax(x)) down every column, computed without forming softmax first: masked entries keep finite log-probabilities instead of log(0) = -Inf, so p * logp stays finite (entropy terms).

Usage

ag_log_softmax(x)

Arguments

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

Details

Mask illegal entries by ADDING a large but FINITE negative number (e.g. -1e4, which also fits f16). -Inf is not supported: the shift below multiplies x by p, and 0 * -Inf is NaN. A column with every entry masked is finite but meaningless (uniform over masked actions); callers must leave at least one entry unmasked.

Stability shift. Any per-column c with max - log(n) <= c <= max keeps exp(x - c) from overflowing and the largest term from underflowing. c = colSums(x * softmax(x)) satisfies it (lse - c is the entropy, at most log(n), and lse >= max), and on the device it is three ggml ops – no max reduction, which ggml lacks along this axis. The result does not depend on c, and the gradient treats it as a constant: dx = g - exp(logp) * colSums(g).

f16 compute dtype: colSums(exp(x - c)) is at most n and is stored in the compute type, so n must stay below 65504 (the f16 maximum); larger vocabularies need ag_dtype("f32"). Inputs of magnitude ~1e4 are themselves quantised to a few units in f16 – a property of the input.

Value

ag_tensor of the same shape.

Examples

lp <- ag_log_softmax(ag_tensor(matrix(c(1, 2, 3, 0, -1e4, 0), 3, 2)))

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