mirror of
https://github.com/ml-explore/mlx.git
synced 2025-08-22 04:56:41 +08:00
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
|
# Copyright © 2023 Apple Inc.
|
||
|
|
|
||
|
|
import itertools
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
import mlx.core as mx
|
||
|
|
import mlx_tests
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
class TestLinalg(mlx_tests.MLXTestCase):
|
||
|
|
def test_norm(self):
|
||
|
|
def check_mx_np(a_mx, a_np):
|
||
|
|
self.assertTrue(np.allclose(a_np, a_mx, atol=1e-5, rtol=1e-6))
|
||
|
|
|
||
|
|
x_mx = mx.arange(18).reshape((2, 3, 3))
|
||
|
|
x_np = np.arange(18).reshape((2, 3, 3))
|
||
|
|
|
||
|
|
for num_axes in range(1, 3):
|
||
|
|
for axis in itertools.combinations(range(3), num_axes):
|
||
|
|
if num_axes == 1:
|
||
|
|
ords = [None, 0.5, 0, 1, 2, 3, -1, 1]
|
||
|
|
else:
|
||
|
|
ords = [None, "fro", -1, 1]
|
||
|
|
for o in ords:
|
||
|
|
for keepdims in [True, False]:
|
||
|
|
if o:
|
||
|
|
out_np = np.linalg.norm(
|
||
|
|
x_np, ord=o, axis=axis, keepdims=keepdims
|
||
|
|
)
|
||
|
|
out_mx = mx.linalg.norm(
|
||
|
|
x_mx, ord=o, axis=axis, keepdims=keepdims
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
out_np = np.linalg.norm(x_np, axis=axis, keepdims=keepdims)
|
||
|
|
out_mx = mx.linalg.norm(x_mx, axis=axis, keepdims=keepdims)
|
||
|
|
assert np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|