(
a: Union[np.ndarray, torch.Tensor],
b: Union[np.ndarray, torch.Tensor],
)
| 28 | # the distance between two N-D tensors given a function. This can be a RMS |
| 29 | # function, maximum abs diff, or any kind of distance function. |
| 30 | def wrapper( |
| 31 | a: Union[np.ndarray, torch.Tensor], |
| 32 | b: Union[np.ndarray, torch.Tensor], |
| 33 | ) -> float: |
| 34 | # convert a and b to np.ndarray type fp64 |
| 35 | a = to_np_arr_fp64(a) |
| 36 | b = to_np_arr_fp64(b) |
| 37 | |
| 38 | # return NaN if shape mismatches |
| 39 | if a.shape != b.shape: |
| 40 | return np.nan |
| 41 | |
| 42 | # After we make sure shape matches, check if it's empty. If yes, return 0 |
| 43 | if a.size == 0: |
| 44 | return 0 |
| 45 | |
| 46 | # np.isinf and np.isnan returns a Boolean mask. Check if Inf or NaN occur at |
| 47 | # the same places in a and b. If not, return NaN |
| 48 | if np.any(np.isinf(a) != np.isinf(b)) or np.any(np.isnan(a) != np.isnan(b)): |
| 49 | return np.nan |
| 50 | |
| 51 | # mask out all the values that are either Inf or NaN |
| 52 | mask = np.isinf(a) | np.isnan(a) |
| 53 | if np.any(mask): |
| 54 | logging.warning("Found inf/nan in tensor when calculating the distance") |
| 55 | |
| 56 | a_masked = a[~mask] |
| 57 | b_masked = b[~mask] |
| 58 | |
| 59 | # after masking, the resulting tensor might be empty. If yes, return 0 |
| 60 | if a_masked.size == 0: |
| 61 | return 0 |
| 62 | |
| 63 | # only compare the rest (those that are actually numbers) using the metric |
| 64 | return fn(a_masked, b_masked) |
| 65 | |
| 66 | return wrapper |
| 67 |
nothing calls this directly
no test coverage detected