Implement custom gradient decorator for graph mode.
(f, *args, **kwargs)
| 211 | |
| 212 | |
| 213 | def _graph_mode_decorator(f, *args, **kwargs): |
| 214 | """Implement custom gradient decorator for graph mode.""" |
| 215 | # TODO(rsepassi): Add support for kwargs |
| 216 | if kwargs: |
| 217 | raise ValueError( |
| 218 | "The custom_gradient decorator currently supports keywords " |
| 219 | "arguments only when eager execution is enabled.") |
| 220 | name = "CustomGradient-%s" % ops.uid() |
| 221 | args = [ops.convert_to_tensor(x) for x in args] |
| 222 | |
| 223 | # Checking global and local variables attempts to ensure that no non-resource |
| 224 | # Variables are added to the graph. |
| 225 | current_var_scope = variable_scope.get_variable_scope() |
| 226 | before_vars = set( |
| 227 | [v.experimental_ref() for v in current_var_scope.global_variables() + |
| 228 | current_var_scope.local_variables()]) |
| 229 | with backprop.GradientTape() as tape: |
| 230 | result, grad_fn = f(*args) |
| 231 | after_vars = set( |
| 232 | [v.experimental_ref() for v in current_var_scope.global_variables() + |
| 233 | current_var_scope.local_variables()]) |
| 234 | new_vars = after_vars - before_vars |
| 235 | new_vars_list = [v.deref() for v in new_vars] |
| 236 | for v in new_vars_list: |
| 237 | if not resource_variable_ops.is_resource_variable(v): |
| 238 | raise TypeError( |
| 239 | "All variables used by a function wrapped with @custom_gradient must " |
| 240 | "be `ResourceVariable`s. Ensure that no `variable_scope` is created " |
| 241 | "with `use_resource=False`.") |
| 242 | # The variables that grad_fn needs to return gradients for are the set of |
| 243 | # variables used that are *not* part of the inputs. |
| 244 | variables_in_tape = frozenset([ |
| 245 | v.experimental_ref() for v in tape.watched_variables() |
| 246 | ]) - frozenset(v.experimental_ref() for v in args) |
| 247 | variables_in_subgraph = frozenset([ |
| 248 | v.experimental_ref() |
| 249 | for v in get_dependent_variables(input_ops=args, output_ops=result) |
| 250 | ]) |
| 251 | variables = list( |
| 252 | [v.deref() for v in variables_in_subgraph.union(variables_in_tape)]) |
| 253 | |
| 254 | grad_argspec = tf_inspect.getfullargspec(grad_fn) |
| 255 | variables_in_signature = ("variables" in grad_argspec.args or |
| 256 | grad_argspec.varkw) |
| 257 | if variables and not variables_in_signature: |
| 258 | raise TypeError("If using @custom_gradient with a function that " |
| 259 | "uses variables, then grad_fn must accept a keyword " |
| 260 | "argument 'variables'.") |
| 261 | if variables_in_signature and not variables: |
| 262 | # User seems to intend to use variables but none were captured. |
| 263 | if not variable_scope.get_variable_scope().use_resource: |
| 264 | raise TypeError("If using @custom_gradient with a function that " |
| 265 | "uses variables, the enclosing variable scope must " |
| 266 | "have use_resource=True.") |
| 267 | else: |
| 268 | logging.warn("@custom_gradient grad_fn has 'variables' in signature, but " |
| 269 | "no ResourceVariables were used on the forward pass.") |
| 270 | flat_result = nest.flatten(result) |
no test coverage detected