def _test_search_sorted_7(test_case, input_dtype, device): sorted_sequence_1d = flow.tensor(np.array([1, 3, 5, 7, 9]), dtype=input_dtype, device=flow.device(device)) gt = np.array(2) output = flow.searchsorted(sorted_sequence_1d, 5, out_int32=True) test_case.assertTrue(np.allclose(output.numpy(), gt, 0.0001, 0.0001)) test_case.assertTrue(output.dtype == flow.int32)
def _test_search_sorted_2(test_case, input_dtype, device): sorted_sequence_1d = flow.tensor(np.array([1, 3, 5, 7, 9]), dtype=input_dtype, device=flow.device(device)) values = flow.tensor(np.array([3, 6, 9]), dtype=input_dtype, device=flow.device(device)) gt = np.array([1, 3, 4]) output = flow.searchsorted(sorted_sequence_1d, values) test_case.assertTrue(np.allclose(output.numpy(), gt, 0.0001, 0.0001)) test_case.assertTrue(output.dtype == flow.int64)
def _test_search_sorted_3(test_case, input_dtype, device): sorted_sequence = flow.tensor( np.array([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]]), dtype=input_dtype, device=flow.device(device), ) values = flow.tensor(np.array([[3, 6, 9], [3, 6, 9]]), dtype=input_dtype, device=flow.device(device)) gt = np.array([[1, 3, 4], [1, 2, 4]]) output = flow.searchsorted(sorted_sequence, values, out_int32=True) test_case.assertTrue(np.allclose(output.numpy(), gt, 0.0001, 0.0001)) test_case.assertTrue(output.dtype == flow.int32)
def _test_search_sorted_4(test_case, input_dtype, device): sorted_sequence = flow.tensor( np.array([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]]), dtype=input_dtype, device=flow.device(device), ) values = flow.tensor(np.array([[3, 6, 9], [3, 6, 9]]), dtype=input_dtype, device=flow.device(device)) sorter = flow.tensor( np.array([[4, 3, 2, 1, 0], [3, 2, 4, 0, 1]]), dtype=flow.int64, device=flow.device(device), ) gt = np.array([[0, 5, 5], [0, 0, 2]]) output = flow.searchsorted(sorted_sequence, values, sorter=sorter) test_case.assertTrue(np.allclose(output.numpy(), gt, 0.0001, 0.0001)) test_case.assertTrue(output.dtype == flow.int64)