Returns whether `input_op` can be used from `op`s context. Conceptually, only inputs from op's while context or any ancestor while context (including outside of any context) are valid. In practice, there are many other edge cases as well. Args: op: Operation input_op: Operation
(op, input_op)
| 264 | |
| 265 | |
| 266 | def CheckInputFromValidContext(op, input_op): |
| 267 | """Returns whether `input_op` can be used from `op`s context. |
| 268 | |
| 269 | Conceptually, only inputs from op's while context or any ancestor while |
| 270 | context (including outside of any context) are valid. In practice, there are |
| 271 | many other edge cases as well. |
| 272 | |
| 273 | Args: |
| 274 | op: Operation |
| 275 | input_op: Operation |
| 276 | |
| 277 | Raises: |
| 278 | ValueError: if input_op is from an invalid context. |
| 279 | """ |
| 280 | op_ctxt = op._get_control_flow_context() # pylint: disable=protected-access |
| 281 | input_ctxt = GetOutputContext(input_op) |
| 282 | valid = False |
| 283 | |
| 284 | if not input_ctxt: |
| 285 | # input_op isn't in a control flow context. |
| 286 | valid = True |
| 287 | elif op_ctxt is input_ctxt: |
| 288 | # input_op is in the same context as op. |
| 289 | valid = True |
| 290 | else: |
| 291 | while_ctxt = GetContainingWhileContext(op_ctxt) |
| 292 | input_while_ctxt = GetContainingWhileContext(input_ctxt) |
| 293 | |
| 294 | if while_ctxt is None: |
| 295 | if input_while_ctxt is None: |
| 296 | # Neither op nor input_op is in a while loop, but one or both are in |
| 297 | # conds. We allow this, although execution will fail if the branch |
| 298 | # corresponding to input_op's cond context isn't taken. |
| 299 | valid = True |
| 300 | # Invalid if op isn't in a while loop and input_op is. Unless... |
| 301 | if IsLoopEnter(op): |
| 302 | # WhileContext._BuildLoop clears context for Enter nodes. |
| 303 | valid = True |
| 304 | if IsSwitch(op): |
| 305 | # CondContext.AddValue clears context for Switch nodes. |
| 306 | valid = True |
| 307 | elif IsContainingContext(while_ctxt, input_while_ctxt): |
| 308 | # input_op is in a while loop which contains op's while loop (or not in a |
| 309 | # while loop at all). |
| 310 | valid = True |
| 311 | elif (while_ctxt.grad_state and |
| 312 | IsContainingContext(while_ctxt.grad_state.forward_context, |
| 313 | input_while_ctxt)): |
| 314 | # op is in a gradient context and input_op is in the associated forward |
| 315 | # pass context or an ancestor thereof. This case is need to build while |
| 316 | # loop gradients. |
| 317 | # NOTE(skyewm): we theoretically also need this case for custom gradient |
| 318 | # functions that close over tensors from ancestor contexts, but I haven't |
| 319 | # verified this. |
| 320 | valid = True |
| 321 | elif (while_ctxt.grad_state and |
| 322 | while_ctxt.grad_state.forward_context is |
| 323 | input_while_ctxt._outer_context): # pylint: disable=protected-access |
nothing calls this directly
no test coverage detected