(self)
| 1130 | self.assertEqual(mx.any(a, axis=1).tolist(), [True, False]) |
| 1131 | |
| 1132 | def test_stop_gradient(self): |
| 1133 | def func(x): |
| 1134 | return mx.sum(2 * x + mx.stop_gradient(3 * x)) |
| 1135 | |
| 1136 | x = mx.array([0.0, 0.1, -3]) |
| 1137 | expected = [2, 2, 2] |
| 1138 | |
| 1139 | self.assertListEqual(mx.grad(func)(x).tolist(), expected) |
| 1140 | |
| 1141 | def test_kron(self): |
| 1142 | # Basic vector test |