ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tf.experimental.unregister_dispatch_for

Unregisters a function that was registered with @dispatch_for_*.

This is primarily intended for testing purposes.

Example:

# Define a type and register a dispatcher to override `tf.abs`:
class MyTensor(tf.experimental.ExtensionType):
  value: tf.Tensor
@dispatch_for_api(tf.abs)
def my_abs(x: MyTensor):
  return MyTensor(tf.abs(x.value))
tf.abs(MyTensor(5))
MyTensor(value=<tf.Tensor: shape=(), dtype=int32, numpy=5>)
# Unregister the dispatcher, so `tf.abs` no longer calls `my_abs`.
unregister_dispatch_for(my_abs)
tf.abs(MyTensor(5))
Traceback (most recent call last):

ValueError: Attempt to convert a value ... to a Tensor.

dispatch_target The function to unregister.

ValueError If dispatch_target was not registered using @dispatch_for, @dispatch_for_unary_elementwise_apis, or @dispatch_for_binary_elementwise_apis.