(self)
| 40 | |
| 41 | @unittest.skipIf(not mx.is_available(mx.gpu), "GPU is not available") |
| 42 | def test_device_context(self): |
| 43 | default = mx.default_device() |
| 44 | diff = mx.cpu if default == mx.gpu else mx.gpu |
| 45 | self.assertNotEqual(default, diff) |
| 46 | with mx.stream(diff): |
| 47 | a = mx.add(mx.zeros((2, 2)), mx.ones((2, 2))) |
| 48 | mx.eval(a) |
| 49 | self.assertEqual(mx.default_device(), diff) |
| 50 | self.assertEqual(mx.default_device(), default) |
| 51 | |
| 52 | def test_op_on_device(self): |
| 53 | x = mx.array(1.0) |