Fix equality check for np arrays with NaNs (#1343)

Summary:
Pull Request resolved: #1343

Prior to this change, having NaN in a field would lead to the equality check failing. This is a real issue with observation data with multi-objective as the non-diagonal entries of covariance matrix get populated with NaNs.

Reviewed By: Balandat

Differential Revision: D42254565

fbshipit-source-id: 8798b89bd268def947e426d0b65859368d51b5a4

ax/utils/common/equality.py

@@ -183,7 +183,7 @@ def object_attribute_dicts_find_unequal_fields(
  list(one_val.values()), list(other_val.values())
  )
  elif isinstance(one_val, np.ndarray):
- equal = np.array_equal(one_val, other_val)
+ equal = np.array_equal(one_val, other_val, equal_nan=True)
  elif isinstance(one_val, datetime):
  equal = datetime_equals(one_val, other_val)
  elif isinstance(one_val, float):

ax/utils/common/tests/test_equality.py

@@ -6,11 +6,13 @@
 
 from datetime import datetime
 
+import numpy as np
 import pandas as pd
 from ax.utils.common.equality import (
  dataframe_equals,
  datetime_equals,
  equality_typechecker,
+ object_attribute_dicts_find_unequal_fields,
  same_elements,
 )
 from ax.utils.common.testutils import TestCase
@@ -46,3 +48,27 @@ def testDataframeEquals(self) -> None:
  self.assertTrue(dataframe_equals(pd.DataFrame(), pd.DataFrame()))
  self.assertTrue(dataframe_equals(pd1, pd2))
  self.assertFalse(dataframe_equals(pd1, pd3))
+
+ def test_numpy_equals(self) -> None:
+ # Simple check.
+ np_0 = {"cov": np.array([[0.1, 0.0], [0.0, 0.1]])}
+ np_1 = {"cov": np.array([[0.1, 0.0], [0.0, 0.1]])}
+ self.assertEqual(
+ object_attribute_dicts_find_unequal_fields(np_0, np_1), ({}, {})
+ )
+ # Unequal.
+ np_1 = {"cov": np.array([[0.1, 0.0], [0.1, 0.1]])}
+ self.assertEqual(
+ object_attribute_dicts_find_unequal_fields(np_0, np_1),
+ ({}, {"cov": (np_0["cov"], np_1["cov"])}),
+ )
+ # With NaNs.
+ np_1 = {"cov": np.array([[0.1, float("nan")], [float("nan"), 0.1]])}
+ self.assertEqual(
+ object_attribute_dicts_find_unequal_fields(np_0, np_1),
+ ({}, {"cov": (np_0["cov"], np_1["cov"])}),
+ )
+ np_0 = {"cov": np.array([[0.1, float("nan")], [float("nan"), 0.1]])}
+ self.assertEqual(
+ object_attribute_dicts_find_unequal_fields(np_0, np_1), ({}, {})
+ )