| 1163 | |
| 1164 | @def_function.function |
| 1165 | def call(c, x): |
| 1166 | |
| 1167 | @custom_gradient.custom_gradient |
| 1168 | def _call(): |
| 1169 | y = c(x) |
| 1170 | |
| 1171 | def grad(dy, variables=None): # pylint: disable=redefined-outer-name |
| 1172 | with backprop.GradientTape(persistent=True) as g: |
| 1173 | g.watch(variables) |
| 1174 | y = c(x) |
| 1175 | grad_vars = [ |
| 1176 | 2 * math_ops.reduce_sum(dy * g.jacobian(y, v)) for v in variables |
| 1177 | ] |
| 1178 | del g |
| 1179 | return (), grad_vars |
| 1180 | |
| 1181 | return y, grad |
| 1182 | |
| 1183 | return _call() |
| 1184 | |
| 1185 | c = MyCallable() |
| 1186 | x = constant_op.constant([1., 2., 3.]) |