View source on GitHub |
Decorator to override default implementation for binary elementwise APIs.
tf.experimental.dispatch_for_binary_elementwise_apis(
x_type, y_type
)
The decorated function (known as the "elementwise api handler") overrides
the default implementation for any binary elementwise API whenever the value
for the first two arguments (typically named x
and y
) match the specified
type annotations. The elementwise api handler is called with two arguments:
elementwise_api_handler(api_func, x, y)
Where x
and y
are the first two arguments to the elementwise api, and
api_func
is a TensorFlow function that takes two parameters and performs the
elementwise operation (e.g., tf.add
).
The following example shows how this decorator can be used to update all
binary elementwise operations to handle a MaskedTensor
type:
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]
Args | |
---|---|
x_type
|
A type annotation indicating when the api handler should be called. |
y_type
|
A type annotation indicating when the api handler should be called. |
Returns | |
---|---|
A decorator. |
Registered APIs
The binary elementwise APIs are:
tf.bitwise.bitwise_and(x, y, name)
tf.bitwise.bitwise_or(x, y, name)
tf.bitwise.bitwise_xor(x, y, name)
tf.bitwise.left_shift(x, y, name)
tf.bitwise.right_shift(x, y, name)
tf.compat.v1.div(x, y, name)
tf.compat.v1.floor_div(x, y, name)
tf.compat.v1.math.scalar_mul(scalar, x, name)
tf.dtypes.complex(real, imag, name)
tf.math.add(x, y, name)
tf.math.atan2(y, x, name)
tf.math.divide(x, y, name)
tf.math.divide_no_nan(x, y, name)
tf.math.equal(x, y, name)
tf.math.floordiv(x, y, name)
tf.math.floormod(x, y, name)
tf.math.greater(x, y, name)
tf.math.greater_equal(x, y, name)
tf.math.less(x, y, name)
tf.math.less_equal(x, y, name)
tf.math.logical_and(x, y, name)
tf.math.logical_or(x, y, name)
tf.math.logical_xor(x, y, name)
tf.math.maximum(x, y, name)
tf.math.minimum(x, y, name)
tf.math.multiply(x, y, name)
tf.math.multiply_no_nan(x, y, name)
tf.math.not_equal(x, y, name)
tf.math.pow(x, y, name)
tf.math.scalar_mul(scalar, x, name)
tf.math.squared_difference(x, y, name)
tf.math.subtract(x, y, name)
tf.math.truediv(x, y, name)
tf.math.xdivy(x, y, name)
tf.math.xlog1py(x, y, name)
tf.math.xlogy(x, y, name)
tf.math.zeta(x, q, name)
tf.nn.sigmoid_cross_entropy_with_logits(labels, logits, name)
tf.realdiv(x, y, name)
tf.truncatediv(x, y, name)
tf.truncatemod(x, y, name)