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

装饰器覆盖一元元素 API 的默认实现。

用法

tf.experimental.dispatch_for_unary_elementwise_apis(
    x_type
)

只要第一个参数(通常命名为 x )的值与类型注释 x_type 匹配，装饰函数(称为“元素 api 处理程序”)就会覆盖任何一元元素 API 的默认实现。 elementwise 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 是：

• tf.bitwise.invert
• tf.cast
• tf.clip_by_value
• tf.compat.v1.math.log_softmax
• tf.compat.v1.ones_like
• tf.compat.v1.strings.length
• tf.compat.v1.strings.substr
• tf.compat.v1.strings.to_hash_bucket(string_tensor, num_buckets, name, input)
• tf.compat.v1.substr
• tf.compat.v1.to_bfloat16
• tf.compat.v1.to_complex128
• tf.compat.v1.to_complex64
• tf.compat.v1.to_double
• tf.compat.v1.to_float
• tf.compat.v1.to_int32
• tf.compat.v1.to_int64
• tf.compat.v1.zeros_like
• tf.debugging.check_numerics
• tf.dtypes.saturate_cast(value, dtype, name)
• tf.image.adjust_brightness
• tf.image.adjust_gamma
• tf.image.convert_image_dtype
• tf.image.random_brightness
• tf.image.stateless_random_brightness
• tf.io.decode_base64(input, name)
• tf.io.decode_compressed(bytes, compression_type, name)
• tf.io.encode_base64(input, pad, name)
• tf.math.abs
• tf.math.acos
• tf.math.acosh
• tf.math.angle
• tf.math.asin
• tf.math.asinh
• tf.math.atan
• tf.math.atanh
• tf.math.bessel_i0
• tf.math.bessel_i0e
• tf.math.bessel_i1
• tf.math.bessel_i1e
• tf.math.ceil
• tf.math.conj
• tf.math.cos
• tf.math.cosh
• tf.math.digamma(x, name)
• tf.math.erf
• tf.math.erfc(x, name)
• tf.math.erfcinv
• tf.math.erfinv(x, name)
• tf.math.exp
• tf.math.expm1
• tf.math.floor
• tf.math.imag
• tf.math.is_finite
• tf.math.is_inf
• tf.math.is_nan
• tf.math.lgamma
• tf.math.log
• tf.math.log1p
• tf.math.log_sigmoid
• tf.math.logical_not
• tf.math.ndtri(x, name)
• tf.math.negative(x, name)
• tf.math.nextafter(x1, x2, name)
• tf.math.real
• tf.math.reciprocal(x, name)
• tf.math.reciprocal_no_nan
• tf.math.rint
• tf.math.round
• tf.math.rsqrt
• tf.math.sigmoid
• tf.math.sign
• tf.math.sin
• tf.math.sinh
• tf.math.softplus
• tf.math.special.bessel_j0
• tf.math.special.bessel_j1
• tf.math.special.bessel_k0
• tf.math.special.bessel_k0e
• tf.math.special.bessel_k1
• tf.math.special.bessel_k1e
• tf.math.special.bessel_y0
• tf.math.special.bessel_y1
• tf.math.special.dawsn
• tf.math.special.expint
• tf.math.special.fresnel_cos
• tf.math.special.fresnel_sin
• tf.math.special.spence
• tf.math.sqrt
• tf.math.square
• tf.math.tan
• tf.math.tanh
• tf.nn.elu
• tf.nn.gelu
• tf.nn.leaky_relu(features, alpha, name)
• tf.nn.relu
• tf.nn.relu6
• tf.nn.selu(features, name)
• tf.nn.silu(features, beta)
• tf.nn.softsign(features, name)
• tf.ones_like
• tf.strings.as_string
• tf.strings.length
• tf.strings.lower
• tf.strings.regex_full_match
• tf.strings.regex_replace
• tf.strings.strip
• tf.strings.substr
• tf.strings.to_hash_bucket
• tf.strings.to_hash_bucket_fast
• tf.strings.to_hash_bucket_strong
• tf.strings.to_number
• tf.strings.unicode_script
• tf.strings.unicode_transcode
• tf.strings.upper
• tf.zeros_like