tf.experimental.dispatch_for_api

Декоратор, переопределяющий реализацию TensorFlow API по умолчанию.

tf.experimental.dispatch_for_api(
    api, *signatures
)

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

Декорированная функция (так называемая «цель диспетчеризации») переопределяет реализацию API по умолчанию, когда API вызывается с параметрами, соответствующими заданной сигнатуре типов. Сигнатуры задаются с помощью словарей, сопоставляющих имена параметров с аннотациями типов. Например, в следующем примере masked_add будет вызвана для tf.add, если и x, и y являются MaskedTensors:

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor

@dispatch_for_api(tf.math.add, {'x': MaskedTensor, 'y': MaskedTensor})
def masked_add(x, y, name=None):
  return MaskedTensor(x.values + y.values, x.mask & y.mask)

mt = tf.add(MaskedTensor([1, 2], [True, False]), MaskedTensor(10, True))
print(f"values={mt.values.numpy()}, mask={mt.mask.numpy()}")
values=[11 12], mask=[ True False]

Если указано несколько сигнатур типов, цель диспетчеризации будет вызвана при совпадении любой из них. Например, следующий код регистрирует masked_add для вызова, если x является MaskedTensor или y является MaskedTensor.

@dispatch_for_api(tf.math.add, {'x': MaskedTensor}, {'y':MaskedTensor})
def masked_add(x, y):
  x_values = x.values if isinstance(x, MaskedTensor) else x
  x_mask = x.mask if isinstance(x, MaskedTensor) else True
  y_values = y.values if isinstance(y, MaskedTensor) else y
  y_mask = y.mask if isinstance(y, MaskedTensor) else True
  return MaskedTensor(x_values + y_values, x_mask & y_mask)

Аннотации типов в сигнатурах типов могут быть объектами типов (например, MaskedTensor), значениями typing.List или значениями typing.Union. Например, следующий код зарегистрирует masked_concat для вызова, если values — это список значений MaskedTensor:

@dispatch_for_api(tf.concat, {'values': typing.List[MaskedTensor]})
def masked_concat(values, axis):
  return MaskedTensor(tf.concat([v.values for v in values], axis),
                      tf.concat([v.mask for v in values], axis))

Каждая сигнатура типа должна содержать хотя бы один подкласс tf.CompositeTensor (включая подклассы tf.ExtensionType), а диспетчеризация запускается, только если хотя бы один параметр с аннотацией типа содержит значение CompositeTensor. Это правило предотвращает вызов диспетчеризации в вырожденных случаях, например в следующих:

• @dispatch_for_api(tf.concat, {'values': List[MaskedTensor]}): При вызове tf.concat([]) пользователем диспетчеризация не будет направлена к декорированной цели диспетчеризации.

• @dispatch_for_api(tf.add, {'x': Union[MaskedTensor, Tensor], 'y': Union[MaskedTensor, Tensor]}): При вызове tf.add(tf.constant(1), tf.constant(2)) пользователем диспетчеризация не будет направлена к декорированной цели диспетчеризации.

Сигнатура цели диспетчеризации должна соответствовать сигнатуре переопределяемого API. В частности, параметры должны иметь те же имена и располагаться в том же порядке. Цель диспетчеризации может не указывать параметр «name»; в этом случае при необходимости она будет обёрнута вызовом tf.name_scope.

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

API TensorFlow, которые можно переопределить с помощью @dispatch_for_api:

<>