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:

<>