Python tf.experimental.dispatch_for_binary_elementwise_apis用法及代码示例

用于覆盖二进制元素 API 的默认实现的装饰器。

用法

tf.experimental.dispatch_for_binary_elementwise_apis(
    x_type, y_type
)

只要前两个参数(通常命名为 x 和 y )的值与指定的类型注释匹配，装饰函数(称为“元素 api 处理程序”)将覆盖任何二进制元素 API 的默认实现。 elementwise api 处理程序使用两个参数调用：

elementwise_api_handler(api_func, x, y)

其中 x 和 y 是 elementwise api 的前两个参数，而 api_func 是一个 TensorFlow 函数，它接受两个参数并执行 elementwise 操作(例如 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 是：

• tf.bitwise.bitwise_and
• tf.bitwise.bitwise_or
• tf.bitwise.bitwise_xor
• tf.bitwise.left_shift
• tf.bitwise.right_shift
• tf.compat.v1.div(x, y, name)
• tf.compat.v1.floor_div(x, y, name)
• tf.compat.v1.scalar_mul
• tf.dtypes.complex
• tf.math.add
• tf.math.atan2
• tf.math.divide
• tf.math.divide_no_nan
• tf.math.equal
• tf.math.floordiv(x, y, name)
• tf.math.floormod(x, y, name)
• tf.math.greater
• tf.math.greater_equal
• tf.math.less
• tf.math.less_equal
• tf.math.logical_and
• tf.math.logical_or
• tf.math.logical_xor
• tf.math.maximum
• tf.math.minimum
• tf.math.multiply
• tf.math.multiply_no_nan(x, y, name)
• tf.math.not_equal
• tf.math.pow
• tf.math.scalar_mul
• tf.math.squared_difference(x, y, name)
• tf.math.subtract
• tf.math.truediv(x, y, name)
• tf.math.xdivy(x, y, name)
• tf.math.xlog1py
• tf.math.xlogy(x, y, name)
• tf.math.zeta(x, q, name)
• tf.nn.sigmoid_cross_entropy_with_logits
• tf.realdiv(x, y, name)
• tf.truncatediv(x, y, name)
• tf.truncatemod(x, y, name)