The gradient of a While op produced by while_loop.
(op, *grads)
| 328 | @ops.RegisterGradient("StatelessWhile") |
| 329 | @ops.RegisterGradient("While") |
| 330 | def _WhileGrad(op, *grads): # pylint: disable=invalid-name |
| 331 | """The gradient of a While op produced by while_loop.""" |
| 332 | # Note that op is not always the same as while_op because the gradient tape, |
| 333 | # for eager mode compatibility, forgets information about the proper op. Since |
| 334 | # the loop cannot run in eager mode, however, we can safely introspect into |
| 335 | # the graph here. |
| 336 | while_op = op.outputs[0].op |
| 337 | cond_graph = _get_graph(while_op, "cond") |
| 338 | body_graph = _get_graph(while_op, "body") |
| 339 | orig_num_params = len(body_graph.outputs) |
| 340 | |
| 341 | maximum_iterations = op.inputs[1] |
| 342 | parallel_iterations = op.get_attr("parallel_iterations") |
| 343 | |
| 344 | try: |
| 345 | num_original_outputs = while_op.get_attr("_num_original_outputs") |
| 346 | except: # pylint: disable=bare-except |
| 347 | num_original_outputs = len(while_op.outputs) |
| 348 | |
| 349 | num_intermediates = len(while_op.outputs) - num_original_outputs |
| 350 | grads = [ |
| 351 | _preprocess_grad(grad, body_out, while_out) # pylint: disable=g-complex-comprehension |
| 352 | for grad, body_out, while_out in zip( |
| 353 | grads[:num_original_outputs], |
| 354 | body_graph.outputs[:num_original_outputs], |
| 355 | while_op.outputs[:num_original_outputs]) |
| 356 | ] + [None] * num_intermediates |
| 357 | |
| 358 | # We compute the gradient for the sub-graph between trainable ys and xs |
| 359 | # with non-None incoming gradients. We later pad the None's to the list of |
| 360 | # outputs. |
| 361 | ys, xs, non_none_grads = zip(*[(y, x, grad) for (y, x, grad) in zip( |
| 362 | body_graph.outputs, body_graph.inputs, grads) if grad is not None]) |
| 363 | |
| 364 | body_grad_graph, args = _create_grad_func( |
| 365 | ys, xs, non_none_grads, cond_graph, body_graph, |
| 366 | util.unique_grad_fn_name(body_graph.name), op, maximum_iterations) |
| 367 | |
| 368 | if body_grad_graph.while_op_needs_rewrite: |
| 369 | # Modify 'op' to output the intermediate accumulators needed by the grad |
| 370 | # function. |
| 371 | # NOTE(skyewm): if there are any active sessions, this modification to `op` |
| 372 | # may make them unrunnable! |
| 373 | |
| 374 | cond_graph.name += "_rewritten" |
| 375 | body_graph.name += "_rewritten" |
| 376 | |
| 377 | new_inputs = body_grad_graph.empty_tensor_lists |
| 378 | new_outputs = body_graph.outputs[orig_num_params:] |
| 379 | |
| 380 | while_op._set_func_attr("cond", util.create_new_tf_function(cond_graph)) |
| 381 | while_op._set_func_attr("body", util.create_new_tf_function(body_graph)) |
| 382 | while_op._set_type_list_attr("T", body_graph.output_types) |
| 383 | while_op._set_shape_list_attr("output_shapes", body_graph.output_shapes) |
| 384 | while_op._add_while_inputs(new_inputs) |
| 385 | while_op._add_outputs([t.dtype for t in new_outputs], |
| 386 | [t.shape for t in new_outputs]) |
| 387 | _copy_handle_data(new_outputs, op.outputs[orig_num_params:]) |
nothing calls this directly
no test coverage detected