Return a function which computes the gradient of `f`.
(f, argnums=0)
| 78 | |
| 79 | |
| 80 | def _grad(f, argnums=0): |
| 81 | """Return a function which computes the gradient of `f`.""" |
| 82 | |
| 83 | def _f(*params): |
| 84 | with backprop.GradientTape() as tape: |
| 85 | tape.watch(params) |
| 86 | primals_out = f(*params) |
| 87 | return tape.gradient( |
| 88 | primals_out, |
| 89 | params[argnums], |
| 90 | unconnected_gradients=UnconnectedGradients.ZERO) |
| 91 | |
| 92 | return _f |
| 93 | |
| 94 | |
| 95 | def _hvp(f, primals, tangents): |
no outgoing calls
no test coverage detected