Reduce mlx arrays

These helpers mirror NumPy-style reductions, optionally collapsing one or more axes. Use drop = FALSE to retain reduced axes with length one (akin to setting drop = FALSE in base R).

Usage

mlx_sum(x, axes = NULL, drop = TRUE)

mlx_prod(x, axes = NULL, drop = TRUE)

mlx_all(x, axes = NULL, drop = TRUE)

mlx_any(x, axes = NULL, drop = TRUE)

mlx_mean(x, axes = NULL, drop = TRUE)

mlx_var(x, axes = NULL, drop = TRUE, ddof = 0L)

mlx_std(x, axes = NULL, drop = TRUE, ddof = 0L)

mlx_sd(x, axes = NULL, drop = TRUE)

Arguments

x

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

axes

Integer vector of axes (1-indexed). Supply positive integers between 1 and the array rank. Many helpers interpret NULL to mean "all axes"—see the function details for specifics.

drop

If TRUE (default), drop dimensions of length 1. If FALSE, retain all dimensions. Equivalent to keepdims = TRUE in underlying mlx functions.

ddof

Non-negative integer delta degrees of freedom for variance or standard deviation reductions.

Value

An mlx array containing the reduction result.

Details

mlx_all() and mlx_any() return mlx boolean scalars, while the base R reducers all() and any() applied to mlx inputs return plain logical scalars.

The axes argument is the inverse of MARGIN in base R apply(). axes gives the axes which will be reduced; MARGIN gives the axes which an operation will be applied over. See the example. Multiple axes are reduced jointly. For variance and standard deviation, ddof is subtracted once from the total number of elements being reduced.

mlx_sd() is a convenience wrapper that matches the default behaviour of stats::sd(), computing a sample standard deviation with ddof = 1.

See also

mlx.core.sum, mlx.core.prod, mlx.core.all, mlx.core.any, mlx.core.mean, mlx.core.var, mlx.core.std

Examples

x <- mlx_matrix(1:4, 2, 2)
mlx_sum(x)
#> mlx array []
#>   dtype: float32
#>   values:
#> [1] 10
mlx_sum(x, axes = 1)
#> mlx array [2]
#>   dtype: float32
#>   values:
#> [1] 3 7
mlx_prod(x, axes = 2, drop = FALSE)
#> mlx array [2 x 1]
#>   dtype: float32
#>   values:
#>      [,1]
#> [1,]    3
#> [2,]    8
mlx_all(x > 0)
#> mlx array []
#>   dtype: bool
#>   values:
#> [1] TRUE
mlx_any(x > 3)
#> mlx array []
#>   dtype: bool
#>   values:
#> [1] TRUE
mlx_mean(x, axes = 1)
#> mlx array [2]
#>   dtype: float32
#>   values:
#> [1] 1.5 3.5
mlx_var(x, axes = 2)
#> mlx array [2]
#>   dtype: float32
#>   values:
#> [1] 1 1
mlx_std(x)
#> mlx array []
#>   dtype: float32
#>   values:
#> [1] 1.118034
mlx_sd(x)
#> mlx array []
#>   dtype: float32
#>   values:
#> [1] 1.290994
# for comparison:
stats::sd(as.matrix(x))
#> [1] 1.290994

a <- array(1:6, dim = 1:3)
ax <- as_mlx(a)
# These are equivalent:
apply(a, 1:2, sum) # leaves dimensions 1-2 intact, sums over dimension 3
#>      [,1] [,2]
#> [1,]    9   12
mlx_sum(a, 3)      # the same
#> mlx array [1 x 2]
#>   dtype: float32
#>   values:
#>      [,1] [,2]
#> [1,]    9   12