(self)
| 984 | self.assertEqual(y.dtype, mx.float32) |
| 985 | |
| 986 | def test_softmax(self): |
| 987 | x = mx.array([1.0, -1.0, 0.0]) |
| 988 | y = nn.softmax(x) |
| 989 | epsilon = 1e-4 |
| 990 | expected_y = mx.array([0.6652, 0.0900, 0.2447]) |
| 991 | self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon)) |
| 992 | self.assertEqual(y.shape, (3,)) |
| 993 | self.assertEqual(y.dtype, mx.float32) |
| 994 | |
| 995 | def test_softmin(self): |
| 996 | x = mx.array([1.0, 2.0, 3.0]) |