| 2810 | ) |
| 2811 | |
| 2812 | def test_bitwise_ops(self): |
| 2813 | types = [ |
| 2814 | mx.uint8, |
| 2815 | mx.uint16, |
| 2816 | mx.uint32, |
| 2817 | mx.uint64, |
| 2818 | mx.int8, |
| 2819 | mx.int16, |
| 2820 | mx.int32, |
| 2821 | mx.int64, |
| 2822 | ] |
| 2823 | a = mx.random.randint(0, 4096, (1000,)) |
| 2824 | b = mx.random.randint(0, 4096, (1000,)) |
| 2825 | for op in ["bitwise_and", "bitwise_or", "bitwise_xor"]: |
| 2826 | for t in types: |
| 2827 | a_mlx = a.astype(t) |
| 2828 | b_mlx = b.astype(t) |
| 2829 | a_np = np.array(a_mlx) |
| 2830 | b_np = np.array(b_mlx) |
| 2831 | out_mlx = getattr(mx, op)(a_mlx, b_mlx) |
| 2832 | out_np = getattr(np, op)(a_np, b_np) |
| 2833 | self.assertTrue(np.array_equal(np.array(out_mlx), out_np)) |
| 2834 | for op in ["left_shift", "right_shift"]: |
| 2835 | for t in types: |
| 2836 | a_mlx = a.astype(t) |
| 2837 | b_mlx = mx.random.randint(0, t.size, (1000,)).astype(t) |
| 2838 | a_np = np.array(a_mlx) |
| 2839 | b_np = np.array(b_mlx) |
| 2840 | out_mlx = getattr(mx, op)(a_mlx, b_mlx) |
| 2841 | out_np = getattr(np, op)(a_np, b_np) |
| 2842 | self.assertTrue(np.array_equal(np.array(out_mlx), out_np)) |
| 2843 | |
| 2844 | for t in types: |
| 2845 | a_mlx = a.astype(t) |
| 2846 | a_np = np.array(a_mlx) |
| 2847 | |
| 2848 | out_mlx = ~a_mlx |
| 2849 | out_np = ~a_np |
| 2850 | self.assertTrue(np.array_equal(np.array(out_mlx), out_np)) |
| 2851 | |
| 2852 | out_mlx = mx.bitwise_invert(a_mlx) |
| 2853 | out_np = mx.bitwise_invert(a_np) |
| 2854 | self.assertTrue(np.array_equal(np.array(out_mlx), out_np)) |
| 2855 | |
| 2856 | # Check broadcasting |
| 2857 | a = mx.ones((3, 1, 5), dtype=mx.bool_) |
| 2858 | b = mx.zeros((1, 2, 5), dtype=mx.bool_) |
| 2859 | c = a | b |
| 2860 | self.assertEqual(c.shape, (3, 2, 5)) |
| 2861 | self.assertTrue(mx.array_equal(c, mx.ones((3, 2, 5), dtype=mx.bool_))) |
| 2862 | |
| 2863 | def test_bitwise_grad(self): |
| 2864 | a = np.random.randint(0, 10, size=(4, 3)) |