Log-sum-exp reduction for mlx arrays

Log-sum-exp reduction for mlx arrays

Usage

mlx_logsumexp(x, axis = NULL, drop = TRUE)

Arguments

x

An mlx array, or an R array/matrix/vector that will be converted via as_mlx().

axis

Optional axis to operate over (1-indexed like R). When NULL, the array is flattened first.

drop

Logical indicating whether the reduced axes should be dropped (default TRUE).

Value

An mlx array containing log-sum-exp results.

See also

mlx.core.logsumexp

Examples

x <- as_mlx(matrix(1:6, 2, 3))
as.matrix(mlx_logsumexp(x))
#> [1] 6.456193
as.matrix(mlx_logsumexp(x, axis = 2))
#> [1] 5.142931 6.142931