tf.searchsorted

Находит позицию, в которую должно быть вставлено значение в отсортированную последовательность.

tf.searchsorted(
    sorted_sequence,
    values,
    side='left',
    out_type=tf.dtypes.int32,
    name=None
)

Этот метод не предназначен для проверки принадлежности (как, например, in в Python).

Обычно эта операция используется для разбиения на интервалы, корзины или дискретизации. Значения values распределяются по индексам корзин на основе границ, указанных в sorted_sequence. Эта операция возвращает индекс корзины для каждого значения.

edges = [-1, 3.3, 9.1, 10.0]
values = [0.0, 4.1, 12.0]
tf.searchsorted(edges, values).numpy()
array([1, 2, 4], dtype=int32)

Аргумент side определяет, какой индекс возвращается, если значение точно совпадает с границей:

seq = [0, 3, 9, 10, 10]
values = [0, 4, 10]
tf.searchsorted(seq, values).numpy()
array([0, 2, 3], dtype=int32)
tf.searchsorted(seq, values, side="right").numpy()
array([1, 2, 5], dtype=int32)

Параметр axis нельзя задать для этой операции. Она всегда выполняется по самой внутренней размерности (axis=-1). Операция может принимать любое количество внешних размерностей. Здесь она применяется к строкам матрицы:

sorted_sequence = [[0., 3., 8., 9., 10.],
                   [1., 2., 3., 4., 5.]]
values = [[9.8, 2.1, 4.3],
          [0.1, 6.6, 4.5, ]]
tf.searchsorted(sorted_sequence, values).numpy()
array([[4, 1, 2],
       [0, 5, 4]], dtype=int32)

Примечание: Предполагается, что sorted_sequence отсортирован вдоль самой внутренней оси, например, с помощью tf.sort(..., axis=-1). Если последовательность не отсортирована, ошибка не возникает, а содержимое возвращаемого тензора не определено.