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

bug fix (#658)

Browse Source
This commit is contained in:
Awni Hannun
2024-02-09 16:50:45 -08:00
committed by GitHub
parent b670485185
commit b96be943dc
4 changed files with 23 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/tests/test_ops.py
View File

@@ -1386,6 +1386,11 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertTrue((a[:-1] < 1e-9).all())
self.assertEqual(a[-1], 1)
# Sliced inputs
y = mx.random.uniform(shape=(8, 4))
out = mx.softmax(y[:, 0:2], axis=-1)
self.assertAlmostEqual(out.sum().item(), 8.0)
def test_concatenate(self):
a_npy = np.random.randn(32, 32, 32)
b_npy = np.random.randn(32, 32, 32)