|
|
|
|
|
|
|
|
import math |
|
|
import unittest |
|
|
from itertools import permutations |
|
|
|
|
|
import mlx.core as mx |
|
|
import mlx_tests |
|
|
import numpy as np |
|
|
|
|
|
try: |
|
|
import torch |
|
|
|
|
|
has_torch = True |
|
|
except ImportError as e: |
|
|
has_torch = False |
|
|
|
|
|
|
|
|
class TestBF16(mlx_tests.MLXTestCase): |
|
|
def __test_ops( |
|
|
self, |
|
|
ref_op, |
|
|
mlx_op, |
|
|
np_args, |
|
|
ref_transform=lambda x: x, |
|
|
mlx_transform=lambda x: mx.array(x), |
|
|
atol=1e-5, |
|
|
): |
|
|
ref_args = map(ref_transform, np_args) |
|
|
mlx_args = map(mlx_transform, np_args) |
|
|
|
|
|
r_ref = ref_op(*ref_args) |
|
|
r_mlx = mlx_op(*mlx_args) |
|
|
|
|
|
self.assertTrue(np.allclose(r_mlx, r_ref, atol=atol)) |
|
|
|
|
|
def __default_test( |
|
|
self, |
|
|
op, |
|
|
np_args, |
|
|
simple_transform=lambda x: x, |
|
|
atol_np=1e-3, |
|
|
atol_torch=1e-5, |
|
|
np_kwargs=dict(), |
|
|
mlx_kwargs=dict(), |
|
|
torch_kwargs=dict(), |
|
|
torch_op=None, |
|
|
): |
|
|
with self.subTest(reference="numpy"): |
|
|
|
|
|
def np_transform(x): |
|
|
x_mx_bf16 = mx.array(x).astype(mx.bfloat16) |
|
|
x_mx_fp32 = x_mx_bf16.astype(mx.float32) |
|
|
return np.asarray(x_mx_fp32) |
|
|
|
|
|
def mlx_fn(*args): |
|
|
out_bf16 = getattr(mx, op)(*args, **mlx_kwargs) |
|
|
return np.asarray(out_bf16.astype(mx.float32)) |
|
|
|
|
|
def np_fn(*args): |
|
|
out_fp32 = getattr(np, op)(*args, **np_kwargs) |
|
|
return np_transform(out_fp32) |
|
|
|
|
|
ref_op = np_fn |
|
|
mlx_op = mlx_fn |
|
|
|
|
|
ref_transform = lambda x: simple_transform(np_transform(x)) |
|
|
mlx_transform = lambda x: simple_transform(mx.array(x).astype(mx.bfloat16)) |
|
|
|
|
|
self.__test_ops( |
|
|
ref_op, |
|
|
mlx_op, |
|
|
np_args, |
|
|
ref_transform=ref_transform, |
|
|
mlx_transform=mlx_transform, |
|
|
atol=atol_np, |
|
|
) |
|
|
|
|
|
if has_torch: |
|
|
with self.subTest(reference="torch"): |
|
|
torch_op = op if torch_op is None else torch_op |
|
|
|
|
|
def torch_fn(*args): |
|
|
out_bf16 = getattr(torch, torch_op)(*args, **torch_kwargs) |
|
|
return out_bf16.to(torch.float32).numpy() |
|
|
|
|
|
ref_op = torch_fn |
|
|
ref_transform = lambda x: simple_transform( |
|
|
torch.from_numpy(x).to(torch.bfloat16) |
|
|
) |
|
|
self.__test_ops( |
|
|
ref_op, |
|
|
mlx_op, |
|
|
np_args, |
|
|
ref_transform=ref_transform, |
|
|
mlx_transform=mlx_transform, |
|
|
atol=atol_torch, |
|
|
) |
|
|
|
|
|
def test_unary_ops(self): |
|
|
x = np.random.rand(18, 28, 38) |
|
|
for op in ["abs", "exp", "log", "square", "sqrt"]: |
|
|
with self.subTest(op=op): |
|
|
np_args = (x.astype(np.float32),) |
|
|
self.__default_test(op, np_args) |
|
|
|
|
|
def test_binary_ops(self): |
|
|
x = np.random.rand(18, 28, 38) |
|
|
y = np.random.rand(18, 28, 38) |
|
|
for op in ["add", "subtract", "multiply", "divide", "maximum", "minimum"]: |
|
|
with self.subTest(op=op): |
|
|
np_args = ( |
|
|
x.astype(np.float32), |
|
|
y.astype(np.float32), |
|
|
) |
|
|
self.__default_test(op, np_args, simple_transform=lambda x: x) |
|
|
self.__default_test(op, np_args, simple_transform=lambda x: x[:1]) |
|
|
self.__default_test(op, np_args, simple_transform=lambda x: x[:, :1]) |
|
|
|
|
|
def test_reduction_ops(self): |
|
|
x = np.random.rand(18, 28, 38).astype(np.float32) |
|
|
|
|
|
for op in ("min", "max"): |
|
|
with self.subTest(op=op): |
|
|
for axes in (0, 1, 2, (0, 1), (0, 2), (1, 2), (0, 1, 2)): |
|
|
with self.subTest(axes=axes): |
|
|
np_args = (x.astype(np.float32),) |
|
|
self.__default_test( |
|
|
op, |
|
|
np_args, |
|
|
np_kwargs={"axis": axes}, |
|
|
mlx_kwargs={"axis": axes}, |
|
|
torch_kwargs={"dim": axes}, |
|
|
torch_op="a" + op, |
|
|
) |
|
|
|
|
|
def test_arg_reduction_ops(self): |
|
|
data = np.random.rand(10, 12, 13).astype(np.float32) |
|
|
x = mx.array(data).astype(mx.bfloat16) |
|
|
data = np.asarray(x.astype(mx.float32)) |
|
|
|
|
|
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) |
|
|
a = a.astype(mx.float32) |
|
|
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) |
|
|
a = a.astype(mx.float32) |
|
|
self.assertEqual(a.tolist(), b.tolist()) |
|
|
a = getattr(mx, op)(x) |
|
|
b = getattr(np, op)(data) |
|
|
a = a.astype(mx.float32) |
|
|
self.assertEqual(a.item(), b) |
|
|
|
|
|
def test_blas_ops(self): |
|
|
if mx.default_device() != mx.gpu: |
|
|
return |
|
|
|
|
|
def test_blas(shape_x, shape_y): |
|
|
np.random.seed(42) |
|
|
with self.subTest(shape_x=shape_x, shape_y=shape_y): |
|
|
x = np.random.normal(0.0, 1.0 / shape_x[-1], size=shape_x) |
|
|
y = np.random.normal(0.0, 1.0 / shape_x[-1], size=shape_y) |
|
|
|
|
|
np_args = ( |
|
|
x.astype(np.float32), |
|
|
y.astype(np.float32), |
|
|
) |
|
|
op = "matmul" |
|
|
|
|
|
self.__default_test(op, np_args, atol_np=1e-3, atol_torch=1e-3) |
|
|
|
|
|
for shape_x, shape_y in [ |
|
|
[(32, 32), (32, 32)], |
|
|
[(23, 57), (57, 1)], |
|
|
[(1, 3), (3, 128)], |
|
|
[(8, 128, 768), (768, 16)], |
|
|
]: |
|
|
test_blas(shape_x, shape_y) |
|
|
|
|
|
@unittest.skipIf(not has_torch, "requires PyTorch") |
|
|
def test_conversion(self): |
|
|
a_torch = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16) |
|
|
a_mx = mx.array(a_torch) |
|
|
expected = mx.array([1.0, 2.0, 3.0], mx.bfloat16) |
|
|
self.assertEqual(a_mx.dtype, mx.bfloat16) |
|
|
self.assertTrue(mx.array_equal(a_mx, expected)) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
mlx_tests.MLXTestRunner() |
|
|
|
