tf.where

Возвращает индексы ненулевых элементов или выбирает элементы из x и y.

tf.where(
    condition, x=None, y=None, name=None
)

Используется в блокнотах

Эта операция работает в двух режимах:

1. Возврат индексов ненулевых элементов — если задан только condition, результатом будет тензор int64, в каждой строке которого указан индекс ненулевого элемента condition. Форма результата — [tf.math.count_nonzero(condition), tf.rank(condition)].
2. Выбор между x и y — если заданы и x, и y, результат имеет форму, полученную в результате широковещательного преобразования x, y и condition. Значение берётся из x, если condition не равно нулю, или из y, если condition равно нулю.

1. Возврат индексов ненулевых элементов

Примечание: В этом режиме тип данных condition может быть bool или любым числовым типом.

Если x и y не заданы (оба имеют значение None):

tf.where вернёт индексы ненулевых элементов condition в виде двумерного тензора формы [n, d], где n — количество ненулевых элементов в condition (tf.count_nonzero(condition)), а d — количество осей condition (tf.rank(condition)).

Индексы выводятся в порядке следования строк. Тип данных condition может быть dtype tf.bool или любым числовым типом dtype.

Здесь condition — тензор bool с одной осью и двумя значениями True. Форма результата — [2,1]

tf.where([True, False, False, True]).numpy()
array([[0],
       [3]])

Здесь condition — целочисленный тензор с двумя осями и тремя ненулевыми значениями. Форма результата — [3, 2].

tf.where([[1, 0, 0], [1, 0, 1]]).numpy()
array([[0, 0],
       [1, 0],
       [1, 2]])

Здесь condition — тензор с плавающей точкой и тремя осями, содержащий пять ненулевых значений. Форма результата — [5, 3].

float_tensor = [[[0.1, 0], [0, 2.2], [3.5, 1e6]],
                [[0,   0], [0,   0], [99,    0]]]
tf.where(float_tensor).numpy()
array([[0, 0, 0],
       [0, 1, 1],
       [0, 2, 0],
       [0, 2, 1],
       [1, 2, 0]])

Эти индексы совпадают с теми, которые tf.sparse.SparseTensor использовал бы для представления тензора условия:

sparse = tf.sparse.from_dense(float_tensor)
sparse.indices.numpy()
array([[0, 0, 0],
       [0, 1, 1],
       [0, 2, 0],
       [0, 2, 1],
       [1, 2, 0]])

Комплексное число считается ненулевым, если ненулевой хотя бы один из его компонентов — действительный или мнимый:

tf.where([complex(0.), complex(1.), 0+1j, 1+1j]).numpy()
array([[1],
       [2],
       [3]])

2. Выбор между x и y

Примечание: В этом режиме тип данных condition должен быть bool.

Если также заданы x и y (оба имеют значения, отличные от None), тензор condition действует как маска, определяющая, из какого тензора следует взять соответствующий элемент или строку результата: из x (если элемент в condition равен True) или из y (если он равен False).

Форма результата определяется широковещательным преобразованием форм condition, x и y.

Если все три входных значения имеют одинаковый размер, операция выполняется поэлементно.

tf.where([True, False, False, True],
         [1, 2, 3, 4],
         [100, 200, 300, 400]).numpy()
array([  1, 200, 300,   4], dtype=int32)

Существуют два основных правила широковещательного преобразования:

1. Если у тензора меньше осей, чем у остальных, слева к его форме добавляются оси длины 1.
2. Оси длины 1 растягиваются до размера соответствующих осей других тензоров.

Вектор длины 1 растягивается до размера остальных векторов:

tf.where([True, False, False, True], [1, 2, 3, 4], [100]).numpy()
array([  1, 100, 100,   4], dtype=int32)

Скаляр расширяется до размера остальных аргументов:

tf.where([[True, False], [False, True]], [[1, 2], [3, 4]], 100).numpy()
array([[  1, 100], [100,   4]], dtype=int32)
tf.where([[True, False], [False, True]], 1, 100).numpy()
array([[  1, 100], [100,   1]], dtype=int32)

Скаляр condition возвращает тензор x или y целиком, применяя широковещательное преобразование.

tf.where(True, [1, 2, 3, 4], 100).numpy()
array([1, 2, 3, 4], dtype=int32)
tf.where(False, [1, 2, 3, 4], 100).numpy()
array([100, 100, 100, 100], dtype=int32)

Рассмотрим более сложный пример широковещательного преобразования: форма condition — [3], форма x — [3,3], а форма y — [3,1]. Сначала широковещательное преобразование расширяет форму condition до [1,3]. Итоговая форма после широковещательного преобразования — [3,3]. condition выберет столбцы из x и y. Поскольку в y только один столбец, все столбцы y будут одинаковыми.

tf.where([True, False, True],
         x=[[1, 2, 3],
            [4, 5, 6],
            [7, 8, 9]],
         y=[[100],
            [200],
            [300]]
).numpy()
array([[ 1, 100, 3],
       [ 4, 200, 6],
       [ 7, 300, 9]], dtype=int32)

Обратите внимание: если градиент любой ветви tf.where порождает NaN, градиент всей операции tf.where будет NaN. Это происходит потому, что при вычислении градиента tf.where две ветви объединяются по соображениям производительности.

Чтобы обойти это ограничение, используйте внутреннюю операцию tf.where, которая гарантирует отсутствие асимптоты у функции, и не вычисляйте значения с градиентом NaN, заменяя опасные входные значения безопасными.

Вместо этого варианта:

x = tf.constant(0., dtype=tf.float32)
with tf.GradientTape() as tape:
  tape.watch(x)
  y = tf.where(x < 1., 0., 1. / x)
print(tape.gradient(y, x))
tf.Tensor(nan, shape=(), dtype=float32)

Хотя значения 1. / x никогда не используются, при x = 0 их градиент равен NaN. Поэтому следует предусмотреть защиту с помощью ещё одной операции tf.where.

x = tf.constant(0., dtype=tf.float32)
with tf.GradientTape() as tape:
  tape.watch(x)
  safe_x = tf.where(tf.equal(x, 0.), 1., x)
  y = tf.where(x < 1., 0., 1. / safe_x)
print(tape.gradient(y, x))
tf.Tensor(0.0, shape=(), dtype=float32)

См. также:

• tf.sparse — индексы, возвращаемые первым вариантом tf.where, могут быть полезны для объектов tf.sparse.SparseTensor.
• tf.gather_nd, tf.scatter_nd и связанные операции — используя список индексов, возвращённый tf.where, можно применять операции семейства scatter и gather для извлечения значений или вставки значений по этим индексам.
• tf.strings.length — тип tf.string не разрешён для condition. Вместо него используйте длину строки.