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

Fix slice data size (#1394)

Browse Source 
* fix slice data size and add tests

* fix contiguous flag

* simplify stride and perform copy for non-contiguous arrays

* fix cpu

* comment
This commit is contained in:
Awni Hannun
2024-09-04 19:10:43 -07:00
committed by GitHub
parent 11371fe251
commit 7cca1727af
12 changed files with 129 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/tests/test_fast.py
View File

@@ -429,6 +429,14 @@ class TestFast(mlx_tests.MLXTestCase):
rx_fast = mx.fast.layer_norm(x, None, None, eps)
self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
def test_slice_into_layer_norm(self):
dim = 128
eps = 1e-5
x = mx.random.uniform(shape=(8, 100, 128))[:, 99:]
rx_fast = mx.fast.layer_norm(x, weight=None, bias=None, eps=eps)
rx = layer_norm(x, None, None, eps)
self.assertLess(mx.abs(rx - rx_fast).max(), 1e-4)
def test_layer_norm_grad(self):
D = 32
eps = 1e-5