Compute quantiles of MLX arrays

Calculates sample quantiles corresponding to given probabilities using linear interpolation (R's type 7 quantiles, the default in stats::quantile()). The S3 method quantile.mlx() provides an interface compatible with the generic stats::quantile().

Usage

mlx_quantile(
  x,
  probs,
  axis = NULL,
  drop = FALSE,
  device = mlx_default_device()
)

# S3 method for class 'mlx'
quantile(x, probs, ...)

Arguments

x

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

probs

Numeric vector of probabilities in [0, 1].

axis

Optional integer axis (or vector of axes) along which to compute quantiles. When NULL (default), quantiles are computed over the entire flattened array.

drop

Logical; when TRUE and computing quantiles along an axis with a single probability, removes the quantile dimension of length 1. Defaults to FALSE to match the behavior of other reduction functions.

device

Execution target: supply "gpu", "cpu", or an mlx_stream created via mlx_new_stream(). Defaults to the current mlx_default_device() unless noted otherwise (helpers that act on an existing array typically reuse that array's device or stream).

...

Additional arguments (currently ignored by quantile.mlx()).

Value

An mlx array containing the requested quantiles. The shape depends on probs, axis, and drop: when axis = NULL, returns a scalar for a single probability or a vector for multiple probabilities. When axis is specified, the quantile dimension replaces the reduced axis (e.g., a (3, 4) matrix with axis = 1 and 2 quantiles gives (2, 4)), unless drop = TRUE with a single probability removes that dimension.

Details

Uses type 7 quantiles (linear interpolation): for probability p and n observations, the quantile is computed as:

• h = (n-1) * p

• Interpolate between floor(h) and ceiling(h)

This matches the default behavior of stats::quantile().

See also

stats::quantile(), mlx.core.sort

Examples

x <- as_mlx(1:10)
mlx_quantile(x, 0.5)  # median
#> mlx array [1]
#>   dtype: float32
#>   device: gpu
#>   values:
#> [1] 5.5
mlx_quantile(x, c(0.25, 0.5, 0.75))  # quartiles
#> mlx array [3]
#>   dtype: float32
#>   device: gpu
#>   values:
#> [1] 3.25 5.50 7.75

# S3 method:
quantile(x, probs = c(0, 0.25, 0.5, 0.75, 1))
#> mlx array [5]
#>   dtype: float32
#>   device: gpu
#>   values:
#> [1]  1.00  3.25  5.50  7.75 10.00

# With axis parameter, quantile dimension replaces the reduced axis:
mat <- mlx_matrix(1:12, 3, 4)  # shape (3, 4)
result <- mlx_quantile(mat, c(0.25, 0.75), axis = 1)  # shape (2, 4)
result <- mlx_quantile(mat, 0.5, axis = 1)  # shape (1, 4)
result <- mlx_quantile(mat, 0.5, axis = 1, drop = TRUE)  # shape (4,)