Create zeros_like for the specified output of an op. If op is in a while loop that is part of gradients(), this method must be called in its grad loop context. Args: op: A tensorflow operation. index: the index for a specific output of the op. Returns: A zero ten
(self, op, index)
| 643 | return result |
| 644 | |
| 645 | def ZerosLike(self, op, index): |
| 646 | """Create zeros_like for the specified output of an op. |
| 647 | |
| 648 | If op is in a while loop that is part of gradients(), this method |
| 649 | must be called in its grad loop context. |
| 650 | |
| 651 | Args: |
| 652 | op: A tensorflow operation. |
| 653 | index: the index for a specific output of the op. |
| 654 | |
| 655 | Returns: |
| 656 | A zero tensor of the same shape of op.outputs[index]. |
| 657 | """ |
| 658 | if util.IsLoopSwitch(op): |
| 659 | return None |
| 660 | if op.graph._building_function: # pylint: disable=protected-access |
| 661 | # The optimization here is tricky to apply to functions |
| 662 | return array_ops.zeros_like(op.outputs[index]) |
| 663 | dead_branch = util.IsSwitch(op) |
| 664 | forward_ctxt = util.GetWhileContext(op) |
| 665 | grad_state = self._map.get(forward_ctxt) |
| 666 | if grad_state is None: |
| 667 | # op is not in a while loop that is part of gradients(). |
| 668 | return ZerosLikeOutsideLoop(op, index) |
| 669 | op_ctxt = op._get_control_flow_context() |
| 670 | val = ops.convert_to_tensor(op.outputs[index], name="tensor") |
| 671 | shape = val.get_shape() |
| 672 | if shape.is_fully_defined(): |
| 673 | # If the shape is known statically, just create a zero tensor with |
| 674 | # the right shape in the grad loop context. |
| 675 | result = constant_op.constant(0, shape=shape.dims, dtype=val.dtype) |
| 676 | if dead_branch: |
| 677 | # op is a cond switch. Guard the zero tensor with a switch. |
| 678 | pred = grad_state.history_map.get(op_ctxt.pred.name) |
| 679 | branch = op_ctxt.branch |
| 680 | result = control_flow_ops._SwitchRefOrTensor(result, pred)[1 - branch] |
| 681 | else: |
| 682 | # Unknown shape so keep a history of the shape at runtime. |
| 683 | if dead_branch: |
| 684 | # Need to add a special switch to guard the value. |
| 685 | pred = op_ctxt.pred |
| 686 | branch = op_ctxt.branch |
| 687 | op_ctxt.outer_context.Enter() |
| 688 | val = control_flow_ops._SwitchRefOrTensor(op.inputs[0], |
| 689 | pred)[1 - branch] |
| 690 | zeros_shape = array_ops.shape_internal(val, optimize=False) |
| 691 | op_ctxt.outer_context.Exit() |
| 692 | val.op._set_control_flow_context(op_ctxt) |
| 693 | zeros_shape.op._set_control_flow_context(op_ctxt) |
| 694 | else: |
| 695 | op_ctxt.Enter() |
| 696 | zeros_shape = array_ops.shape_internal(val, optimize=False) |
| 697 | op_ctxt.Exit() |
| 698 | |
| 699 | # Add forward accumulator for shape. |
| 700 | grad_state.grad_context.Exit() |
| 701 | history_zeros_shape = grad_state.AddForwardAccumulator( |
| 702 | zeros_shape, dead_branch=dead_branch) |
no test coverage detected