Inner function closure for calculating gradients.
(*args, **kwargs)
| 385 | |
| 386 | @custom_gradient |
| 387 | def inner(*args, **kwargs): |
| 388 | """Inner function closure for calculating gradients.""" |
| 389 | result = f(*args, **kwargs) |
| 390 | |
| 391 | def grad(dresult, variables=None): |
| 392 | """Gradient function calculation for inner function.""" |
| 393 | with backprop.GradientTape() as t: |
| 394 | t.watch(args) |
| 395 | if variables is not None: |
| 396 | t.watch(variables) |
| 397 | with ops.control_dependencies([dresult]): |
| 398 | result = f(*args, **kwargs) |
| 399 | kw_vars = [] |
| 400 | if variables is not None: |
| 401 | kw_vars = list(variables) |
| 402 | grads = t.gradient( |
| 403 | result, list(args) + kw_vars, output_gradients=[dresult]) |
| 404 | return grads[:len(args)], grads[len(args):] |
| 405 | |
| 406 | return result, grad |
| 407 | |
| 408 | return inner |
| 409 |