|
|
|
|
|
|
|
|
import unittest |
|
|
from itertools import combinations, permutations |
|
|
|
|
|
import mlx.core as mx |
|
|
import mlx_tests |
|
|
import numpy as np |
|
|
|
|
|
|
|
|
class TestReduce(mlx_tests.MLXTestCase): |
|
|
def test_axis_permutation_sums(self): |
|
|
for shape in [(5, 5, 1, 5, 5), (65, 65, 1, 65)]: |
|
|
with self.subTest(shape=shape): |
|
|
x_npy = (np.random.randn(*shape) * 128).astype(np.int32) |
|
|
x_mlx = mx.array(x_npy) |
|
|
for t in permutations(range(len(shape))): |
|
|
with self.subTest(t=t): |
|
|
y_npy = np.transpose(x_npy, t) |
|
|
y_mlx = mx.transpose(x_mlx, t) |
|
|
for n in range(1, len(shape) + 1): |
|
|
for a in combinations(range(len(shape)), n): |
|
|
with self.subTest(a=a): |
|
|
z_npy = np.sum(y_npy, axis=a) |
|
|
z_mlx = mx.sum(y_mlx, axis=a) |
|
|
mx.eval(z_mlx) |
|
|
self.assertTrue(np.all(z_npy == z_mlx)) |
|
|
|
|
|
def test_expand_sums(self): |
|
|
x_npy = np.random.randn(5, 1, 5, 1, 5, 1).astype(np.float32) |
|
|
x_mlx = mx.array(x_npy) |
|
|
for m in range(1, 4): |
|
|
for ax in combinations([1, 3, 5], m): |
|
|
shape = np.array([5, 1, 5, 1, 5, 1]) |
|
|
shape[list(ax)] = 5 |
|
|
shape = shape.tolist() |
|
|
with self.subTest(shape=shape): |
|
|
y_npy = np.broadcast_to(x_npy, shape) |
|
|
y_mlx = mx.broadcast_to(x_mlx, shape) |
|
|
for n in range(1, 7): |
|
|
for a in combinations(range(6), n): |
|
|
with self.subTest(a=a): |
|
|
z_npy = np.sum(y_npy, axis=a) / 1000 |
|
|
z_mlx = mx.sum(y_mlx, axis=a) / 1000 |
|
|
mx.eval(z_mlx) |
|
|
self.assertTrue( |
|
|
np.allclose(z_npy, np.array(z_mlx), atol=1e-4) |
|
|
) |
|
|
|
|
|
def test_dtypes(self): |
|
|
int_dtypes = [ |
|
|
"int8", |
|
|
"int16", |
|
|
"int32", |
|
|
"uint8", |
|
|
"uint16", |
|
|
"uint32", |
|
|
"int64", |
|
|
"uint64", |
|
|
"complex64", |
|
|
] |
|
|
float_dtypes = ["float32"] |
|
|
|
|
|
for dtype in int_dtypes + float_dtypes: |
|
|
with self.subTest(dtype=dtype): |
|
|
x = np.random.uniform(0, 2, size=(3, 3, 3)).astype(getattr(np, dtype)) |
|
|
y = mx.array(x) |
|
|
|
|
|
for op in ("sum", "prod", "min", "max"): |
|
|
with self.subTest(op=op): |
|
|
np_op = getattr(np, op) |
|
|
mlx_op = getattr(mx, op) |
|
|
|
|
|
for axes in (None, 0, 1, 2, (0, 1), (0, 2), (1, 2), (0, 1, 2)): |
|
|
with self.subTest(axes=axes): |
|
|
if op in ("sum", "prod"): |
|
|
r_np = np_op( |
|
|
x, axis=axes, dtype=(getattr(np, dtype)) |
|
|
) |
|
|
else: |
|
|
r_np = np_op(x, axis=axes) |
|
|
r_mlx = mlx_op(y, axis=axes) |
|
|
mx.eval(r_mlx) |
|
|
self.assertTrue(np.allclose(r_np, r_mlx, atol=1e-4)) |
|
|
|
|
|
def test_arg_reduce(self): |
|
|
dtypes = [ |
|
|
"uint8", |
|
|
"uint16", |
|
|
"uint32", |
|
|
"uint64", |
|
|
"int8", |
|
|
"int16", |
|
|
"int32", |
|
|
"int64", |
|
|
"float16", |
|
|
"float32", |
|
|
] |
|
|
for dtype in dtypes: |
|
|
with self.subTest(dtype=dtype): |
|
|
data = np.random.rand(10, 12, 13).astype(getattr(np, dtype)) |
|
|
x = mx.array(data) |
|
|
for op in ["argmin", "argmax"]: |
|
|
for axis in range(3): |
|
|
for kd in [True, False]: |
|
|
a = getattr(mx, op)(x, axis, kd) |
|
|
b = getattr(np, op)(data, axis, keepdims=kd) |
|
|
self.assertEqual(a.tolist(), b.tolist()) |
|
|
|
|
|
for op in ["argmin", "argmax"]: |
|
|
a = getattr(mx, op)(x, keepdims=True) |
|
|
b = getattr(np, op)(data, keepdims=True) |
|
|
self.assertEqual(a.tolist(), b.tolist()) |
|
|
a = getattr(mx, op)(x) |
|
|
b = getattr(np, op)(data) |
|
|
self.assertEqual(a.item(), b) |
|
|
|
|
|
def test_edge_case(self): |
|
|
x = (mx.random.normal((100, 1, 100, 100)) * 128).astype(mx.int32) |
|
|
x = x.transpose(0, 3, 1, 2) |
|
|
|
|
|
y = x.sum((0, 2, 3)) |
|
|
mx.eval(y) |
|
|
z = np.array(x).sum((0, 2, 3)) |
|
|
self.assertTrue(np.all(z == y)) |
|
|
|
|
|
def test_sum_bool(self): |
|
|
x = np.random.uniform(0, 1, size=(10, 10, 10)) > 0.5 |
|
|
y = mx.array(x) |
|
|
npsum = x.sum().item() |
|
|
mxsum = y.sum().item() |
|
|
self.assertEqual(npsum, mxsum) |
|
|
|
|
|
def test_many_reduction_axes(self): |
|
|
|
|
|
def check(x, axes): |
|
|
expected = x |
|
|
for ax in axes: |
|
|
expected = mx.sum(expected, axis=ax, keepdims=True) |
|
|
out = mx.sum(x, axis=axes, keepdims=True) |
|
|
self.assertTrue(mx.array_equal(out, expected)) |
|
|
|
|
|
x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4)) |
|
|
check(x, (0, 2, 4)) |
|
|
|
|
|
x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4)) |
|
|
check(x, (0, 2, 4, 6)) |
|
|
|
|
|
x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4, 4, 4)) |
|
|
check(x, (0, 2, 4, 6, 8)) |
|
|
|
|
|
x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4, 4, 4, 128)) |
|
|
x = x.transpose(1, 0, 2, 3, 4, 5, 6, 7, 8, 9) |
|
|
check(x, (1, 3, 5, 7, 9)) |
|
|
|
|
|
def test_nan_propagation(self): |
|
|
dtypes = [ |
|
|
"uint8", |
|
|
"uint16", |
|
|
"uint32", |
|
|
"int8", |
|
|
"int16", |
|
|
"int32", |
|
|
"float16", |
|
|
"float32", |
|
|
] |
|
|
|
|
|
for dtype in dtypes: |
|
|
with self.subTest(dtype=dtype): |
|
|
x = (mx.random.normal((4, 4)) * 10).astype(getattr(mx, dtype)) |
|
|
indices = mx.random.randint(0, 4, shape=(6,)).reshape(3, 2) |
|
|
for idx in indices: |
|
|
x[idx[0], idx[1]] = mx.nan |
|
|
x_np = np.array(x) |
|
|
|
|
|
for op in ["max", "min"]: |
|
|
for axis in [0, 1]: |
|
|
out = getattr(mx, op)(x, axis=axis) |
|
|
ref = getattr(np, op)(x_np, axis=axis) |
|
|
self.assertTrue(np.array_equal(out, ref, equal_nan=True)) |
|
|
|
|
|
def test_nan_propagation_complex64(self): |
|
|
complex_array_1 = mx.array( |
|
|
[1 + 1j, 2 + 2j, 3 + 3j, mx.nan + 4j], dtype=mx.complex64 |
|
|
).reshape(2, 2) |
|
|
complex_array_2 = mx.array( |
|
|
[1 + 1j, 2 + 2j, 3 + mx.nan * 1j, 4 + 4j], dtype=mx.complex64 |
|
|
).reshape(2, 2) |
|
|
complex_array_3 = mx.array( |
|
|
[1 + 1j, 2 + mx.nan * 1j, 3 + 3j, 4 + 4j], dtype=mx.complex64 |
|
|
).reshape(2, 2) |
|
|
complex_array_4 = mx.array( |
|
|
[mx.nan + 1j, 2 + 2j, 3 + 3j, 4 + 4j], dtype=mx.complex64 |
|
|
).reshape(2, 2) |
|
|
|
|
|
np_arrays = [ |
|
|
np.array(complex_array_1), |
|
|
np.array(complex_array_2), |
|
|
np.array(complex_array_3), |
|
|
np.array(complex_array_4), |
|
|
] |
|
|
|
|
|
for mx_arr, np_arr in zip( |
|
|
[complex_array_1, complex_array_2, complex_array_3, complex_array_4], |
|
|
np_arrays, |
|
|
): |
|
|
for axis in [0, 1]: |
|
|
for op in ["max", "min"]: |
|
|
out = getattr(mx, op)(mx_arr, axis=axis) |
|
|
ref = getattr(np, op)(np_arr, axis=axis) |
|
|
self.assertTrue(np.array_equal(out, ref, equal_nan=True)) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
mlx_tests.MLXTestRunner(failfast=True) |
|
|
|
