Custom grad fn wrapper.
(*result_grads)
| 272 | all_tensors = flat_result + args + variables |
| 273 | |
| 274 | def tape_grad_fn(*result_grads): |
| 275 | """Custom grad fn wrapper.""" |
| 276 | result_grads = result_grads[:flat_result_len] |
| 277 | if variables: |
| 278 | input_grads, variable_grads = grad_fn(*result_grads, variables=variables) |
| 279 | if len(variable_grads) != len(variables): |
| 280 | raise ValueError("Must return gradient for each variable from " |
| 281 | "@custom_gradient grad_fn.") |
| 282 | else: |
| 283 | input_grads = grad_fn(*result_grads) |
| 284 | variable_grads = [] |
| 285 | |
| 286 | # Need to return one value per input to the IdentityN, so pad the |
| 287 | # gradients of the inputs of the custom_gradient function with the |
| 288 | # gradients of the outputs as well. |
| 289 | input_grads = nest.flatten(input_grads) |
| 290 | return ([None] * flat_result_len) + input_grads + variable_grads |
| 291 | |
| 292 | @ops.RegisterGradient(name) |
| 293 | def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable |
no test coverage detected