| 549 | |
| 550 | |
| 551 | class _Loop(NetBuilder): |
| 552 | def __init__(self, iters=None, name=None): |
| 553 | NetBuilder.__init__(self, name, _stop_blob_required=True) |
| 554 | if iters is not None: |
| 555 | self._inc = ops.Const(1) |
| 556 | self._iter = ops.Const(0) |
| 557 | self._num_iters = ( |
| 558 | iters if isinstance(iters, core.BlobReference) |
| 559 | else ops.Const(iters)) |
| 560 | else: |
| 561 | self._num_iters = None |
| 562 | |
| 563 | def iter(self): |
| 564 | assert self._num_iters is not None, ( |
| 565 | 'This loop does not have a number of iterations.') |
| 566 | assert self._iter is not None, ( |
| 567 | 'iter() must be called from inside the loop context') |
| 568 | return self._iter |
| 569 | |
| 570 | def __enter__(self): |
| 571 | builder = NetBuilder.__enter__(self) |
| 572 | if self._num_iters is not None: |
| 573 | ops.stop_if(ops.GE([self._iter, self._num_iters])) |
| 574 | return builder |
| 575 | |
| 576 | def __exit__(self, type, *args): |
| 577 | if type is None and self._num_iters is not None: |
| 578 | self.current_net().Add([self._iter, self._inc], [self._iter]) |
| 579 | NetBuilder.__exit__(self, type, *args) |
| 580 | |
| 581 | |
| 582 | class _RunIf(_RunOnce): |
no outgoing calls
no test coverage detected
searching dependent graphs…