| 1798 | self.assertTrue(np.array_equal(c, c_)) |
| 1799 | |
| 1800 | def test_softmax(self): |
| 1801 | cases = [(np.float32, 1e-6), (np.float16, 1e-3)] |
| 1802 | |
| 1803 | for dtype, atol in cases: |
| 1804 | a_npy = np.random.randn(16, 8, 32).astype(dtype) |
| 1805 | a_mlx = mx.array(a_npy) |
| 1806 | |
| 1807 | def np_softmax(x, axis): |
| 1808 | ex = np.exp(x - np.max(x, axis=axis, keepdims=True)) |
| 1809 | return ex / np.sum(ex, axis=axis, keepdims=True) |
| 1810 | |
| 1811 | for axes in (None, 0, 1, 2, (0, 1), (1, 2), (0, 2), (0, 1, 2)): |
| 1812 | b_npy = np_softmax(a_npy, axes) |
| 1813 | b_mlx = mx.softmax(a_mlx, axes) |
| 1814 | self.assertTrue(np.allclose(b_npy, b_mlx, atol=atol)) |
| 1815 | |
| 1816 | for s in [100, 2049, 4097, 8193]: |
| 1817 | a = np.full(s, -np.inf) |
| 1818 | a[-1] = 0.0 |
| 1819 | a = mx.softmax(mx.array(a)) |
| 1820 | self.assertFalse(np.any(np.isnan(a))) |
| 1821 | self.assertTrue((a[:-1] < 1e-9).all()) |
| 1822 | self.assertEqual(a[-1], 1) |
| 1823 | |
| 1824 | # Sliced inputs |
| 1825 | y = mx.random.uniform(shape=(8, 4)) |
| 1826 | out = mx.softmax(y[:, 0:2], axis=-1) |
| 1827 | self.assertAlmostEqual(out.sum().item(), 8.0, 5) |
| 1828 | |
| 1829 | # Precise |
| 1830 | for t in [mx.float16, mx.bfloat16]: |
| 1831 | a = (10 * mx.random.normal(shape=(1024,))).astype(t) |
| 1832 | out_expect = mx.softmax(a.astype(mx.float32)).astype(t) |
| 1833 | out = mx.softmax(a, axis=-1, precise=True) |
| 1834 | self.assertTrue(mx.allclose(out_expect, out)) |
| 1835 | |
| 1836 | # All Infs give NaNs |
| 1837 | for n in [127, 128, 129]: |
| 1838 | x = mx.full((n,), vals=-float("inf")) |
| 1839 | self.assertTrue(mx.all(mx.isnan(mx.softmax(x)))) |
| 1840 | |
| 1841 | # Transposed inputs |
| 1842 | a = mx.random.uniform(shape=(32, 32, 32)) |
| 1843 | b = mx.softmax(a, axis=-1) |
| 1844 | c = mx.softmax(a.swapaxes(0, 1), axis=-1).swapaxes(0, 1) |
| 1845 | self.assertEqual((b - c).abs().max().item(), 0.0) |
| 1846 | |
| 1847 | with self.assertRaises(ValueError): |
| 1848 | mx.softmax(mx.array(1.0), axis=-1) |
| 1849 | |
| 1850 | def test_concatenate(self): |
| 1851 | a_npy = np.random.randn(32, 32, 32) |