| 2075 | self.assertTrue(np.array_equal(a, anp)) |
| 2076 | |
| 2077 | def test_setitem_with_boolean_mask(self): |
| 2078 | # Python list mask |
| 2079 | a = mx.array([1.0, 2.0, 3.0]) |
| 2080 | mask = [True, False, True] |
| 2081 | src = mx.array([5.0, 6.0]) |
| 2082 | expected = mx.array([5.0, 2.0, 6.0]) |
| 2083 | a[mask] = src |
| 2084 | self.assertTrue(mx.array_equal(a, expected)) |
| 2085 | |
| 2086 | # mx.array scalar mask |
| 2087 | a = mx.array([1.0, 2.0, 3.0]) |
| 2088 | mask = mx.array(True) |
| 2089 | expected = mx.array([5.0, 5.0, 5.0]) |
| 2090 | a[mask] = 5.0 |
| 2091 | self.assertTrue(mx.array_equal(a, expected)) |
| 2092 | |
| 2093 | # scalar mask |
| 2094 | a = mx.array([1.0, 2.0, 3.0]) |
| 2095 | mask = True |
| 2096 | expected = mx.array([5.0, 5.0, 5.0]) |
| 2097 | a[mask] = 5.0 |
| 2098 | self.assertTrue(mx.array_equal(a, expected)) |
| 2099 | |
| 2100 | mask_np = np.zeros((1, 10, 10), dtype=bool) |
| 2101 | with self.assertRaises(ValueError): |
| 2102 | mx.arange(1000).reshape(10, 10, 10)[mask_np] = 0 |
| 2103 | |
| 2104 | mask_np = np.zeros((10, 10, 1), dtype=bool) |
| 2105 | with self.assertRaises(ValueError): |
| 2106 | mx.arange(1000).reshape(10, 10, 10)[mask_np] = 0 |
| 2107 | |
| 2108 | def test_array_namespace(self): |
| 2109 | a = mx.array(1.0) |