| 3153 | self.assertTrue(mx.array_equal(a, mx.arange(8))) |
| 3154 | |
| 3155 | def test_complex_ops(self): |
| 3156 | x = mx.array( |
| 3157 | [ |
| 3158 | 3.0 + 4.0j, |
| 3159 | -5.0 + 12.0j, |
| 3160 | -8.0 + 0.0j, |
| 3161 | 0.0 + 9.0j, |
| 3162 | 0.0 + 0.0j, |
| 3163 | ] |
| 3164 | ) |
| 3165 | |
| 3166 | ops = ["arccos", "arcsin", "arctan", "square", "sqrt"] |
| 3167 | for op in ops: |
| 3168 | with self.subTest(op=op): |
| 3169 | np_op = getattr(np, op) |
| 3170 | mx_op = getattr(mx, op) |
| 3171 | self.assertTrue(np.allclose(mx_op(x), np_op(x))) |
| 3172 | |
| 3173 | x = mx.array( |
| 3174 | [ |
| 3175 | 3.0 + 4.0j, |
| 3176 | -5.0 + 12.0j, |
| 3177 | -8.0 + 0.0j, |
| 3178 | 0.0 + 9.0j, |
| 3179 | 9.0 + 1.0j, |
| 3180 | ] |
| 3181 | ) |
| 3182 | self.assertTrue(np.allclose(mx.rsqrt(x), 1.0 / np.sqrt(x))) |
| 3183 | |
| 3184 | def test_complex_power(self): |
| 3185 | out = mx.power(mx.array(0j), 2) |