Implement custom gradient decorator for eager mode.
(f, *args, **kwargs)
| 313 | |
| 314 | |
| 315 | def _eager_mode_decorator(f, *args, **kwargs): |
| 316 | """Implement custom gradient decorator for eager mode.""" |
| 317 | with backprop.GradientTape() as tape: |
| 318 | result, grad_fn = f(*args, **kwargs) |
| 319 | all_inputs = list(args) + list(kwargs.values()) |
| 320 | # The variables that grad_fn needs to return gradients for are the set of |
| 321 | # variables used that are *not* part of the inputs. |
| 322 | variables = [ |
| 323 | v.deref() # pylint: disable=g-complex-comprehension |
| 324 | for v in set(v.experimental_ref() for v in tape.watched_variables()) |
| 325 | if all(v.deref() is not i for i in all_inputs) |
| 326 | ] |
| 327 | grad_argspec = tf_inspect.getfullargspec(grad_fn) |
| 328 | if (variables and ("variables" not in grad_argspec.args) and |
| 329 | not grad_argspec.varkw): |
| 330 | raise TypeError("If using @custom_gradient with a function that " |
| 331 | "uses variables, then grad_fn must accept a keyword " |
| 332 | "argument 'variables'.") |
| 333 | flat_result = nest.flatten(result) |
| 334 | # TODO(apassos) consider removing the identity below. |
| 335 | flat_result = [gen_array_ops.identity(x) for x in flat_result] |
| 336 | |
| 337 | input_tensors = [ops.convert_to_tensor(x) for x |
| 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) |
| 359 | flat_result = list(flat_result) |
| 360 | return nest.pack_sequence_as(result, flat_result) |
| 361 | |
| 362 | |
| 363 | @tf_export("recompute_grad") |
no test coverage detected