Custom grad fn wrapper.
(*result_grads)
| 338 | in list(args) + list(variables)] |
| 339 | arg_count = len(args) |
| 340 | def actual_grad_fn(*result_grads): |
| 341 | """Custom grad fn wrapper.""" |
| 342 | if variables: |
| 343 | input_grads, variable_grads = grad_fn(*result_grads, variables=variables) |
| 344 | if len(variable_grads) != len(variables): |
| 345 | raise ValueError("Must return gradient for each variable from " |
| 346 | "@custom_gradient grad_fn.") |
| 347 | else: |
| 348 | input_grads = grad_fn(*result_grads) |
| 349 | variable_grads = [] |
| 350 | flat_grads = nest.flatten(input_grads) |
| 351 | if len(flat_grads) != arg_count: |
| 352 | raise ValueError( |
| 353 | "custom_gradient function expected to return", arg_count, |
| 354 | "gradients but returned", len(flat_grads), "instead.") |
| 355 | return nest.flatten(input_grads) + variable_grads |
| 356 | |
| 357 | tape_lib.record_operation(f.__name__, flat_result, input_tensors, |
| 358 | actual_grad_fn) |