jax-ml / jax-ml/ml_dtypes

NaN comparison fails in np.testing.assert_equal

Open
#301 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
361
Forks
60
Avg merge
18h 41m
Merged PRs (30d)
3

Description

The test utility `np.testing.assert_equal` will treat NaN as equal values. However this is not the case for some ml_dtypes arrays:

```py
import ml_dtypes
import numpy as np

# This will succeed
fp32_array = np.array(np.nan, dtype=np.float32)
np.testing.assert_equal(fp32_array, fp32_array)

# This will fail
array = np.array(np.nan, dtype=ml_dtypes.bfloat16)
np.testing.assert_equal(array, array)
```

with

```pytb
Traceback (most recent call last):
File "", line 1, in
File "numpy/testing/_private/utils.py", line 371, in assert_equal
return assert_array_equal(actual, desired, err_msg, verbose,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "numpy/testing/_private/utils.py", line 1051, in assert_array_equal
assert_array_compare(operator.__eq__, actual, desired, err_msg=err_msg,
File "numpy/testing/_private/utils.py", line 916, in assert_array_compare
raise AssertionError(msg)
AssertionError:
Arrays are not equal

Mismatched elements: 1 / 1 (100%)
Max absolute difference among violations: nan
Max relative difference among violations: nan
ACTUAL: array(nan, dtype=bfloat16)
DESIRED: array(nan, dtype=bfloat16)
```

Contributor guide

Open the contributing guide

Research direction

Start by running the provided Python reproduction with NumPy and ml_dtypes, then inspect the assert_equal and assert_array_compare paths in numpy/testing/_private/utils.py named by the traceback. Determine where bfloat16 NaN comparisons diverge from float32 behavior. Done means the reproduction succeeds and a regression test covers the bfloat16 NaN case.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
testing-qa
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.