tf.experimental.dispatch_for_unary_elementwise_apis

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

tf.experimental.dispatch_for_unary_elementwise_apis(
    x_type
)

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

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

elementwise_api_handler(api_func, x)

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

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

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor
@dispatch_for_unary_elementwise_apis(MaskedTensor)
def unary_elementwise_api_handler(api_func, x):
  return MaskedTensor(api_func(x.values), x.mask)
mt = MaskedTensor([1, -2, -3], [True, False, True])
abs_mt = tf.abs(mt)
print(f"values={abs_mt.values.numpy()}, mask={abs_mt.mask.numpy()}")
values=[1 2 3], mask=[ True False True]

Для унарных элементных операций, которые принимают дополнительные аргументы помимо x, эти аргументы не передаются обработчику элементных API, но автоматически добавляются при вызове api_func. Например, в следующем примере параметр dtype не передается в unary_elementwise_api_handler, но добавляется функцией api_func.

ones_mt = tf.ones_like(mt, dtype=tf.float32)
print(f"values={ones_mt.values.numpy()}, mask={ones_mt.mask.numpy()}")
values=[1.0 1.0 1.0], mask=[ True False True]

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

Унарные элементные API:

<>