(
a: mx.array,
b: mx.array,
precision: mx.Dtype,
rtol: float | None = None,
atol: float | None = None,
message: str | None = None,
)
| 9 | |
| 10 | |
| 11 | def assert_allclose( |
| 12 | a: mx.array, |
| 13 | b: mx.array, |
| 14 | precision: mx.Dtype, |
| 15 | rtol: float | None = None, |
| 16 | atol: float | None = None, |
| 17 | message: str | None = None, |
| 18 | ): |
| 19 | if a.dtype == mx.bfloat16: |
| 20 | a = a.astype(mx.float32) |
| 21 | if b.dtype == mx.bfloat16: |
| 22 | b = b.astype(mx.float32) |
| 23 | a = np.array(a) |
| 24 | b = np.array(b) |
| 25 | if precision == mx.float32: |
| 26 | rtol = rtol or 1.0e-5 |
| 27 | atol = atol or 1.0e-6 |
| 28 | elif precision == mx.float16: |
| 29 | rtol = rtol or 5.0e-2 |
| 30 | atol = atol or 1.0e-3 |
| 31 | elif precision == mx.bfloat16: |
| 32 | rtol = rtol or 5.0e-2 |
| 33 | atol = atol or 1.0e-2 |
| 34 | else: |
| 35 | raise ValueError(f"Unsupported precision: {precision}") |
| 36 | assert a.shape == b.shape, f"shape mismatch: {a.shape} vs {b.shape}" |
| 37 | if not np.allclose(a, b, rtol=rtol, atol=atol): |
| 38 | diff = np.invert(np.isclose(a, b, rtol=rtol, atol=atol)) |
| 39 | with np.printoptions(precision=3, suppress=True): |
| 40 | print("a=", a) |
| 41 | print("b=", b) |
| 42 | print("diff_a=", a * diff) |
| 43 | print("diff_b=", b * diff) |
| 44 | print("diff_a_val=", a[diff]) |
| 45 | print("diff_b_val=", b[diff]) |
| 46 | assert False, f"result mismatch: {message}" |
| 47 | |
| 48 | |
| 49 | def np_type_to_mx_type(np_type: np.dtype) -> mx.Dtype: |
nothing calls this directly
no outgoing calls
no test coverage detected