Context to colocate with `op` if `colocate_gradients_with_ops`.
(op, gradient_uid, colocate_gradients_with_ops)
| 301 | |
| 302 | @contextlib.contextmanager |
| 303 | def _maybe_colocate_with(op, gradient_uid, colocate_gradients_with_ops): # pylint: disable=invalid-name |
| 304 | """Context to colocate with `op` if `colocate_gradients_with_ops`.""" |
| 305 | if colocate_gradients_with_ops: |
| 306 | with ops._colocate_with_for_gradient(op, gradient_uid): # pylint: disable=protected-access |
| 307 | yield |
| 308 | else: |
| 309 | yield |
| 310 | |
| 311 | |
| 312 | def _IsPartitionedCall(op): |
no test coverage detected