jax-ml / jax-ml/ml_dtypes

Support np.testing.assert_array_max_ulp

Open
#56 1 comment 1 reaction 0 assignees View on GitHub
contributions welcome enhancement
Dominant language
C++
Stars
361
Forks
60
Avg merge
18h 41m
Merged PRs (30d)
3

Description

When working with bfloat16, given how variable the precision is, it would be super useful to be able to make assertions based on the number of ULPs between numbers rather than traditional measures of tolerance like `atol`/`rtol`.

However, `np.testing.assert_array_max_ulp` does not currently work with the `bfloat16` type from `ml_dtypes`:

```python
>>> a = np.array(0.0, dtype=ml_dtypes.bfloat16)
>>> np.testing.assert_array_max_ulp(a, a)
Traceback (most recent call last):
File "", line 1, in
File "numpy/testing/_private/utils.py", line 1638, in assert_array_max_ulp
ret = nulp_diff(a, b, dtype)
File "numpy/testing/_private/utils.py", line 1686, in nulp_diff
t = np.common_type(x, y)
File "<__array_function__ internals>", line 180, in common_type
File "numpy/lib/type_check.py", line 730, in common_type
raise TypeError("can't get common type for non-numeric array")
TypeError: can't get common type for non-numeric array
```

Given how `ml_dtypes` integrates with NumPy/the existence of NumPy's own `float16` dtype, I'm not sure whether this is possible - but I'll open this issue on the off chance that it might be :)

Thanks!

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the bfloat16 example and inspect NumPy's numpy/testing/_private/utils.py, especially assert_array_max_ulp, nulp_diff, and the np.common_type call named in the traceback. Done means the assertion accepts ml_dtypes.bfloat16 and has coverage for the reported case without breaking existing behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, testing-qa
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.