tf.experimental.dispatch_for_binary_elementwise_apis

Декоратор для переопределения стандартной реализации бинарных элементных API.

tf.experimental.dispatch_for_binary_elementwise_apis(
    x_type, y_type
)

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

Декорированная функция (известная как "обработчик элементных API") переопределяет стандартную реализацию любого бинарного элементного API, когда значения для первых двух аргументов (обычно именованных x и y) соответствуют указанным аннотациям типов. Обработчик элементных API вызывается с двумя аргументами:

elementwise_api_handler(api_func, x, y)

Где x и y — первые два аргумента элементного API, а api_func — функция TensorFlow, которая принимает два параметра и выполняет элементную операцию (например, tf.add).

Следующий пример показывает, как этот декоратор можно использовать для обновления всех бинарных элементных операций для обработки типа MaskedTensor:

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor
@dispatch_for_binary_elementwise_apis(MaskedTensor, MaskedTensor)
def binary_elementwise_api_handler(api_func, x, y):
  return MaskedTensor(api_func(x.values, y.values), x.mask & y.mask)
a = MaskedTensor([1, 2, 3, 4, 5], [True, True, True, True, False])
b = MaskedTensor([2, 4, 6, 8, 0], [True, True, True, False, True])
c = tf.add(a, b)
print(f"values={c.values.numpy()}, mask={c.mask.numpy()}")
values=[ 3 6 9 12 5], mask=[ True True True False False]

Зарегистрированные API

Бинарные элементные API:

<>