Elementwise conditional selection
Usage
mlx_where(condition, x, y)
Arguments
- condition
Logical mlx array (non-zero values are treated as TRUE).
- x, y
Arrays broadcastable to the shape of condition.
Value
An mlx array where elements are drawn from x when
condition is TRUE, otherwise from y.
Details
Behaves like ifelse() for arrays, but evaluates both branches.
Examples
cond <- as_mlx(matrix(c(TRUE, FALSE, TRUE, FALSE), 2, 2))
a <- as_mlx(matrix(1:4, 2, 2))
b <- as_mlx(matrix(5:8, 2, 2))
mlx_where(cond, a, b)
#> mlx array [2 x 2]
#> dtype: float32
#> device: gpu
#> values:
#> [,1] [,2]
#> [1,] 1 3
#> [2,] 6 8