Add `op` to the current context.
(self, op)
| 1691 | return result |
| 1692 | |
| 1693 | def AddOp(self, op): |
| 1694 | """Add `op` to the current context.""" |
| 1695 | # For a reduction op, if op is in a grad context and its input is from |
| 1696 | # its forward context, moving op to the forward context means we would |
| 1697 | # store the tensor after the reduction as opposed to the tensor before |
| 1698 | # reduction, and therefore could significantly reduce memory consumption. |
| 1699 | # For now, we do this only for a few ops. |
| 1700 | # |
| 1701 | # If in XLA context, do not move constant ops to forward pass as pushing to |
| 1702 | # and popping from a stack removes the constant property of an op and breaks |
| 1703 | # XLA compilation, which requires certain inputs to be constant for certain |
| 1704 | # ops. |
| 1705 | if not util.IsInXLAContext(op) and op.type in {"Shape", "Size", "Rank"}: |
| 1706 | grad_ctxt = ops.get_default_graph()._get_control_flow_context() |
| 1707 | if grad_ctxt: |
| 1708 | grad_ctxt = grad_ctxt.GetWhileContext() |
| 1709 | if grad_ctxt.grad_state: |
| 1710 | op_input_forward_ctxt = util.GetWhileContext(op.inputs[0].op) |
| 1711 | if op_input_forward_ctxt == grad_ctxt.grad_state.forward_context: |
| 1712 | op_input_ctxt = op.inputs[0].op._get_control_flow_context() |
| 1713 | op._set_control_flow_context(op_input_ctxt) |
| 1714 | op_input_ctxt._AddOpInternal(op) |
| 1715 | return |
| 1716 | self._AddOpInternal(op) |
| 1717 | |
| 1718 | def _AddOpInternal(self, op): |
| 1719 | """Add `op` to the current context. |
nothing calls this directly
no test coverage detected