numpy function imports

mellon/validation.py

@@ -1,7 +1,7 @@
 from collections.abc import Iterable
 import logging
 
-from jax.numpy import asarray, concatenate, isscalar, full, ndarray, where
+from jax.numpy import asarray, concatenate, isscalar, full, ndarray, where, isnan, isinf
 from jax.numpy import sum as arraysum
 from jax.numpy import min as arraymin
 from jax.numpy import all as arrayall
@@ -477,8 +477,8 @@ def _validate_nn_distances(nn_distances):
 
     zeros = nn_distances == 0
     # Check for invalid values
-    nan_count = np.isnan(nn_distances).sum()
-    inf_count = np.isinf(nn_distances).sum()
+    nan_count = isnan(nn_distances).sum()
+    inf_count = isinf(nn_distances).sum()
     negative_count = (nn_distances < 0).sum()
     if nan_count > 0 or inf_count > 0 or negative_count > 0:
         total_invalid = nan_count + inf_count + negative_count