zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 16:48:10 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

added tri / tril / triu (#170)

Browse Source 
* added tri / tril / triu

* fixed tests

* ctest tests

* tri overload and simplified tests

* changes from comment

* more tests for m

* ensure assert if not 2-D

* remove broadcast_to

* minor tweaks

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Diogo
2023-12-15 20:30:34 -05:00
committed by GitHub
parent 2e02acdc83
commit dc2edc762c
9 changed files with 207 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/tests/test_nn.py
View File

@@ -451,15 +451,13 @@ class TestNN(mlx_tests.MLXTestCase):
def test_prelu(self):
self.assertEqualArray(
[mx.array([1.0, -1.0, 0.0, 0.5])],
nn.PReLU(),
nn.PReLU()(mx.array([1.0, -1.0, 0.0, 0.5])),
mx.array([1.0, -0.25, 0.0, 0.5]),
)
def test_mish(self):
self.assertEqualArray(
[mx.array([1.0, -1.0, 0.0, 0.5])],
nn.Mish(),
nn.Mish()(mx.array([1.0, -1.0, 0.0, 0.5])),
mx.array([0.8651, -0.3034, 0.0000, 0.3752]),
)