Create zeros_like for the specified output of an op.
(op, index)
| 775 | |
| 776 | |
| 777 | def ZerosLikeOutsideLoop(op, index): |
| 778 | """Create zeros_like for the specified output of an op.""" |
| 779 | val = op.outputs[index] |
| 780 | if not util.IsSwitch(op): |
| 781 | if val.dtype == dtypes.resource: |
| 782 | return array_ops.zeros( |
| 783 | gen_resource_variable_ops.variable_shape(val), |
| 784 | dtype=default_gradient.get_zeros_dtype(val)) |
| 785 | return array_ops.zeros_like(val, optimize=False) |
| 786 | else: |
| 787 | op_ctxt = op._get_control_flow_context() |
| 788 | if op_ctxt: |
| 789 | # We are in a cond context. Use a switch to create zeros only when needed. |
| 790 | pred = op_ctxt.pred |
| 791 | branch = op_ctxt.branch |
| 792 | switch_val = control_flow_ops.switch(op.inputs[0], pred)[1 - branch] |
| 793 | # A op is created along the branch taken as control dependencies are on |
| 794 | # the whole op and not on the tensor output. |
| 795 | pivot = array_ops.identity(switch_val) |
| 796 | if val.dtype == dtypes.resource: |
| 797 | with ops.control_dependencies([pivot]): |
| 798 | return array_ops.zeros( |
| 799 | gen_resource_variable_ops.variable_shape(switch_val), |
| 800 | dtype=default_gradient.get_zeros_dtype(val)) |
| 801 | zeros_shape = array_ops.shape_internal(switch_val, optimize=False) |
| 802 | # Ensure ops created within array_ops.zeros are dominated by switch in |
| 803 | # cond context. |
| 804 | with ops.control_dependencies([pivot]): |
| 805 | return array_ops.zeros(zeros_shape, dtype=val.dtype) |
| 806 | else: |
| 807 | return array_ops.zeros_like(val, optimize=False) |
no test coverage detected