Reorder mlx array axes

mlx_moveaxis() repositions one or more axes to new locations.
aperm.mlx() provides the familiar R interface, permuting axes according to perm via repeated calls to mlx_moveaxis().

Usage

mlx_moveaxis(x, source, destination)

# S3 method for class 'mlx'
aperm(a, perm = NULL, resize = TRUE, ...)

x, a: An object coercible to mlx via as_mlx().
source: Integer vector of axis indices to move (1-indexed; negatives count from the end).
destination: Integer vector giving the target positions for source axes (1-indexed; negatives count from the end). Must be the same length as source.
perm: Integer permutation describing the desired axis order, matching the semantics of base::aperm().
resize: Logical flag from base::aperm(). Only TRUE is currently supported for mlx arrays.
...: Additional arguments accepted for compatibility; ignored.

An mlx array with axes permuted.

x <- as_mlx(array(1:8, dim = c(2, 2, 2)))
moved <- mlx_moveaxis(x, source = 1, destination = 3)
permuted <- aperm(x, c(2, 1, 3))