def searchsorted(a: NdarrayOrTensor, v: NdarrayOrTensor, right=False, sorter=None): side = "right" if right else "left" if isinstance(a, np.ndarray): return np.searchsorted(a, v, side, sorter) # type: ignore if hasattr(torch, "searchsorted"): return torch.searchsorted(a, v, right=right) # type: ignore # if using old PyTorch, will convert to numpy array then compute ret = np.searchsorted(a.cpu().numpy(), v.cpu().numpy(), side, sorter) # type: ignore ret, *_ = convert_to_dst_type(ret, a) return ret
def assert_allclose( actual: NdarrayOrTensor, desired: NdarrayOrTensor, type_test: bool = True, device_test: bool = False, *args, **kwargs, ): """ Assert that types and all values of two data objects are close. Args: actual: Pytorch Tensor or numpy array for comparison. desired: Pytorch Tensor or numpy array to compare against. type_test: whether to test that `actual` and `desired` are both numpy arrays or torch tensors. device_test: whether to test the device property. args: extra arguments to pass on to `np.testing.assert_allclose`. kwargs: extra arguments to pass on to `np.testing.assert_allclose`. """ if type_test: # check both actual and desired are of the same type np.testing.assert_equal(isinstance(actual, np.ndarray), isinstance(desired, np.ndarray), "numpy type") np.testing.assert_equal(isinstance(actual, torch.Tensor), isinstance(desired, torch.Tensor), "torch type") if isinstance(desired, torch.Tensor) or isinstance(actual, torch.Tensor): if device_test: np.testing.assert_equal(str(actual.device), str(desired.device), "torch device check") # type: ignore actual = actual.cpu().numpy() if isinstance(actual, torch.Tensor) else actual desired = desired.cpu().numpy() if isinstance( desired, torch.Tensor) else desired np.testing.assert_allclose(actual, desired, *args, **kwargs)