tf.keras.ops.where

Return elements chosen from x1 or x2 depending on condition.

conditionWhere True, yield x1, otherwise yield x2.
x1Values from which to choose when condition is True.
x2Values from which to choose when condition is False.

A tensor with elements from x1 where condition is True, and elements from x2 where condition is False.