View source on GitHub |
A decorator for registering the gradient function for an op type.
tf.RegisterGradient( op_type )
This decorator is only used when defining a new op type. For an op with m
inputs and n
outputs, the gradient function is a function that takes the original Operation
and n
Tensor
objects (representing the gradients with respect to each output of the op), and returns m
Tensor
objects (representing the partial gradients with respect to each input of the op).
For example, assuming that operations of type "Sub"
take two inputs x
and y
, and return a single output x - y
, the following gradient function would be registered:
@tf.RegisterGradient("Sub") def _sub_grad(unused_op, grad): return grad, tf.negative(grad)
The decorator argument op_type
is the string type of an operation. This corresponds to the OpDef.name
field for the proto that defines the operation.
Args | |
---|---|
op_type | The string type of an operation. This corresponds to the OpDef.name field for the proto that defines the operation. |
Raises | |
---|---|
TypeError | If op_type is not string. |
__call__
__call__( f )
Registers the function f
as gradient function for op_type
.
© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.3/api_docs/python/tf/RegisterGradient