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)
