zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-04 15:04:40 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add inner / outer op (#348)

Browse Source 
* inner / outer impl

* python tests

* ops list and ack

* updated descriptions

* use test helper

* removed dtype check and flatten outer to 1-D

* updated docs

* just use the reshape to flatten
This commit is contained in:
Diogo
2024-01-07 12:01:09 -05:00
committed by GitHub
parent 6ea6b4258d
commit 449b43762e
7 changed files with 140 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
python/tests/test_ops.py
View File

@@ -1547,6 +1547,34 @@ class TestOps(mlx_tests.MLXTestCase):
dims=([2, 1, 3], [1, 2, 0]),
)
def test_inner(self):
self.assertCmpNumpy([(3,), (3,)], mx.inner, np.inner)
self.assertCmpNumpy([(1, 1, 2), (3, 2)], mx.inner, np.inner)
self.assertCmpNumpy([(2, 3, 4), (4,)], mx.inner, np.inner)
def test_outer(self):
self.assertCmpNumpy([(3,), (3,)], mx.outer, np.outer)
self.assertCmpNumpy(
[
mx.ones(
5,
),
mx.linspace(-2, 2, 5),
],
mx.outer,
np.outer,
)
self.assertCmpNumpy(
[
1j * mx.linspace(2, -2, 5),
mx.ones(
5,
),
],
mx.outer,
np.outer,
)
if __name__ == "__main__":
unittest.main()