Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/autograph/operators/control_flow.py: 12%

1# Copyright 2018 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15"""Control flow statements: loops, conditionals, etc. 

16 

17Note: most of these operators accept pairs of get_state/set_state functions, to 

18capture mutations that the corresponding code blocks might make. These 

19mutations only need to be captured when staging the control flow, and they just 

20work when reverting to Python behavior. 

21 

22__Examples__ 

23 

24``` 

25while cond: 

26 self.x += i 

27``` 

28 

29When the functionalized version is executed as a Python loop, it just works: 

30 

31``` 

32def loop_body(): 

33 self.x += i # works as expected for Python loops 

34``` 

35 

36But it won't work for TF loops: 

37 

38``` 

39def loop_body(): 

40 self.x += i # self.x has the wrong value! 

41``` 

42 

43get_state/set_state allow piping the mutations through the loop variables as 

44well, in effect changing the loop body: 

45 

46``` 

47def loop_body(self_x): 

48 self.x = self_x # self.x now has the proper value 

49 self.x += i # the original block 

50 self_x = self.x # write self.x back into the loop vars 

51 return self_x 

52 

53self_x = tf.while_loop(...) 

54self.x = self_x # the result is not properly captured 

55``` 

56""" 

57 

58import functools 

59import sys 

60import traceback 

61 

62import numpy as np 

63 

64from tensorflow.python.autograph.operators import py_builtins 

65from tensorflow.python.autograph.operators import variables 

66from tensorflow.python.autograph.utils import ag_logging 

67from tensorflow.python.autograph.utils import misc 

68from tensorflow.python.autograph.utils import tensors 

69from tensorflow.python.autograph.utils import type_registry 

70from tensorflow.python.framework import dtypes 

71from tensorflow.python.framework import errors_impl 

72from tensorflow.python.framework import func_graph 

73from tensorflow.python.framework import ops 

74from tensorflow.python.framework import tensor_conversion 

75from tensorflow.python.framework import tensor_shape 

76from tensorflow.python.framework import tensor_util 

77from tensorflow.python.ops import array_ops 

78from tensorflow.python.ops import cond as tf_cond 

79from tensorflow.python.ops import control_flow_assert 

80from tensorflow.python.ops import control_flow_util 

81from tensorflow.python.ops import math_ops 

82from tensorflow.python.ops import tensor_array_ops 

83from tensorflow.python.ops import while_loop 

84from tensorflow.python.ops.ragged import ragged_tensor 

85from tensorflow.python.types import distribute 

86from tensorflow.python.util import nest 

87from tensorflow.python.util import variable_utils 

88 

89 

90PYTHON_MAX_ITERATIONS = 100000000 # Fails in about one minute for empty loops. 

91WARN_INEFFICIENT_UNROLL = True 

92INEFFICIENT_UNROLL_MIN_ITERATIONS = 50000 

93INEFFICIENT_UNROLL_MIN_OPS = 1 

94 

95 

96# TODO(mdan): Use the custom operator pattern instead of type dispatch. 

97# An example of this pattern is found in the implementation of distributed 

98# datasets. Before it can be used though, we need to standardize the interface. 

99 

100for_loop_registry = type_registry.TypeRegistry() 

101 

102 

103def _is_none_or_undef(value): 

104 """Tests whether a value is None or undefined. 

105 

106 AutoGraph represents undefined symbols using special objects of type Undefined 

107 or UndefinedReturnValue. 

108 

109 Args: 

110 value: value to test 

111 

112 Returns: 

113 Boolean 

114 """ 

115 return ((value is None) 

116 or isinstance(value, variables.UndefinedReturnValue) 

117 or isinstance(value, variables.Undefined)) 

118 

119 

120def _verify_tf_condition(cond, tag): 

121 """Ensures that the condition can be used in a TF control flow.""" 

122 extra_hint = 'to check for None, use `is not None`' 

123 cond = tensor_conversion.convert_to_tensor_v2(cond) 

124 

125 if cond.dtype != dtypes.bool: 

126 raise ValueError( 

127 'condition of {} expected to be `tf.bool` scalar, got {}' 

128 '; to use as boolean Tensor, use `tf.cast`' 

129 '; {}'.format(tag, cond, extra_hint)) 

130 

131 if cond.shape is None or cond.shape.ndims is None: 

132 # TODO(mdan): Consider a explicit size check, if not too slow. 

133 cond = array_ops.reshape(cond, ()) 

134 

135 elif cond.shape.ndims > 0: 

136 known_dims = [d for d in cond.shape.as_list() if d is not None] 

137 if np.prod(known_dims) > 1: 

138 raise ValueError( 

139 'condition of {} expected to be `tf.bool` scalar, got {}' 

140 '; {}'.format(tag, cond, extra_hint)) 

141 else: 

142 cond = array_ops.reshape(cond, ()) 

143 

144 return cond 

145 

146 

147def verify_loop_init_vars( 

148 init_vars, symbol_names, first_iter_vars=None, extra_message=None 

149): 

150 """Ensures that all values in the state are valid to use in a TF loop. 

151 

152 The init_vars may contain placeholder values derived from first_iter_vars. 

153 

154 Args: 

155 init_vars: initial loop variables (as taken before entering the loop) 

156 symbol_names: corresponding names of the initial loop variables 

157 first_iter_vars: loop variables after one iteration of the loop 

158 extra_message: an extra string to append to the error message, in case of 

159 "undefined variable" errors (see variables.Undefined) 

160 """ 

161 if not symbol_names: 

162 return 

163 if first_iter_vars is None: 

164 first_iter_vars = (None,) * len(symbol_names) 

165 

166 assert len(symbol_names) == len(init_vars) 

167 assert len(symbol_names) == len(first_iter_vars) 

168 for name, val, fi_val in zip(symbol_names, init_vars, first_iter_vars): 

169 if isinstance(val, variables.UndefinedReturnValue): 

170 if fi_val: 

171 raise ValueError( 

172 'the return value from a TensorFlow loop may only be a {}; got {}' 

173 .format(LEGAL_LOOP_TYPES, type(fi_val))) 

174 else: 

175 # TODO(mdan): This can be handled by removing the return value. 

176 raise NotImplementedError( 

177 'a return statement cannot be placed inside this TensorFlow loop;' 

178 ' this may happen if a return statement depends on a' 

179 ' static Python condition such as a hyperparameter') 

180 

181 error_msg = None 

182 if val is None: 

183 error_msg = "'{}' is not allowed to be None before the loop".format(name) 

184 elif isinstance(val, variables.Undefined): 

185 error_msg = "'{}' must be defined before the loop".format(name) 

186 if extra_message: 

187 error_msg += '\n' + extra_message 

188 

189 if error_msg is not None: 

190 raise ValueError(error_msg) 

191 

192 

193def _is_subshape(left, right): 

194 """Returns True if left shape is at least as specific as right shape.""" 

195 # TODO(mdan): This code should be in TensorShape. 

196 # Note: this is not the same as TensorShape.is_compatible_with, which is 

197 # symmetric. 

198 # This code also duplicates _ShapeLessThanOrEqual from control_flow_ops.py. 

199 if right.dims is None: 

200 return True 

201 if left.ndims != right.ndims: 

202 return False 

203 for ldim, rdim in zip(left.dims, right.dims): 

204 if rdim.value is not None and ldim.value != rdim.value: 

205 return False 

206 return True 

207 

208 

209# TODO(mdan): Remove these verifications once TF ops can properly report names. 

210def _verify_single_loop_var( 

211 name, check_shape, init, entry, exit_, shape_invariant): 

212 """Verifies whether the initial, entry and exit values are consistent.""" 

213 assert entry is not None, "no TF op should set '{}' to None?".format(name) 

214 if exit_ is None: 

215 raise ValueError("'{}' is None at the end of the iteration.".format(name)) 

216 

217 if isinstance(init, (bool, int, float, str, np.ndarray)): 

218 init = tensor_conversion.convert_to_tensor_v2(init) 

219 if isinstance(entry, (bool, int, float, str, np.ndarray)): 

220 entry = tensor_conversion.convert_to_tensor_v2(entry) 

221 if isinstance(exit_, (bool, int, float, str, np.ndarray)): 

222 exit_ = tensor_conversion.convert_to_tensor_v2(exit_) 

223 

224 if (not tensor_util.is_tf_type(entry) or 

225 not tensor_util.is_tf_type(exit_)): 

226 return 

227 

228 # TODO(mdan): Properly account for CompositeTensors. 

229 if (not hasattr(entry, 'dtype') or 

230 not hasattr(exit_, 'dtype')): 

231 return 

232 if (not hasattr(entry, 'shape') or 

233 not hasattr(exit_, 'shape')): 

234 return 

235 

236 if entry.dtype != exit_.dtype: 

237 raise TypeError( 

238 "'{}' has dtype {} before the loop, but dtype {} after one" 

239 ' iteration'.format( 

240 name, 

241 entry.dtype.name, 

242 exit_.dtype.name, 

243 )) 

244 if check_shape: 

245 exit_shape = exit_.shape 

246 if shape_invariant is None: 

247 entry_shape = entry.shape 

248 if not _is_subshape(exit_shape, entry_shape): 

249 raise ValueError( 

250 "'{}' has shape {} before the loop, but shape {} after one" 

251 ' iteration. Use tf.autograph.experimental.set_loop_options to set' 

252 ' shape invariants.'.format(name, entry_shape, exit_shape)) 

253 else: 

254 init_shape = init.shape 

255 if not _is_subshape(init_shape, shape_invariant): 

256 raise ValueError( 

257 "'{}' has shape {} before the loop, which does not conform with" 

258 ' the shape invariant {}.'.format(name, init_shape, 

259 shape_invariant)) 

260 if not _is_subshape(exit_shape, shape_invariant): 

261 raise ValueError( 

262 "'{}' has shape {} after one iteration, which does not conform with" 

263 ' the shape invariant {}.'.format(name, exit_shape, shape_invariant) 

264 ) 

265 

266 

267def verify_tf_loop_vars( 

268 init_vars, 

269 iter_entry_vars, 

270 iter_exit_vars, 

271 symbol_names, 

272 opts, 

273 check_shapes=True, 

274): 

275 """Verifies loop variables for consistency.""" 

276 if check_shapes and 'shape_invariants' in opts: 

277 shape_invariants = opts['shape_invariants'] 

278 else: 

279 shape_invariants = nest.map_structure(lambda _: None, iter_entry_vars) 

280 

281 assert len(symbol_names) == len(shape_invariants) 

282 assert len(symbol_names) == len(init_vars) 

283 assert len(symbol_names) == len(iter_entry_vars) 

284 assert len(symbol_names) == len(iter_exit_vars) 

285 

286 for i in range(len(symbol_names)): 

287 name = symbol_names[i] 

288 init = init_vars[i] 

289 entry = iter_entry_vars[i] 

290 exit_ = iter_exit_vars[i] 

291 invariant = shape_invariants[i] 

292 

293 try: 

294 nest.assert_same_structure(init, entry, expand_composites=True) 

295 except (ValueError, TypeError): 

296 # `Variable`s in `init` may be implicitly converted to `Tensor`s. Convert 

297 # `ResourceVariable`s to Tensors so tf.nest.assert_same_structure 

298 # won't break due to type spec mismatches between `ResourceVariable`s and 

299 # `Tensor`s. 

300 try: 

301 init_tensors = variable_utils.convert_variables_to_tensors(init) 

302 nest.assert_same_structure(init_tensors, entry, expand_composites=True) 

303 except (ValueError, TypeError) as e: 

304 raise TypeError("'{}' does not have the same nested structure after one" 

305 ' iteration.\n\n{}'.format(name, e)) from e 

306 

307 try: 

308 nest.assert_same_structure(entry, exit_, expand_composites=True) 

309 except (ValueError, TypeError) as e: 

310 raise TypeError("'{}' does not have the same nested structure after one" 

311 ' iteration.\n\n{}'.format(name, e)) from e 

312 if invariant is not None: 

313 try: 

314 nest.assert_same_structure(init, invariant, expand_composites=False) 

315 except (ValueError, TypeError) as e: 

316 raise TypeError("'{}' does not have the same nested structure as its" 

317 ' corresponding shape invariant.\n\n{}'.format( 

318 name, e)) from e 

319 

320 nest.map_structure( 

321 functools.partial(_verify_single_loop_var, name, check_shapes), init, 

322 entry, exit_, invariant) 

323 

324 

325def verify_single_cond_var(name, body_var, orelse_var): 

326 """Verifies whether body_var and orelse_var are consistent.""" 

327 if body_var is None: 

328 raise ValueError("'{}' is None at the end of the main branch.".format(name)) 

329 if orelse_var is None: 

330 raise ValueError( 

331 "'{}' is None at the end of the else branch.".format(name)) 

332 

333 if isinstance(body_var, (bool, int, float, str, np.ndarray)): 

334 body_var = tensor_conversion.convert_to_tensor_v2(body_var) 

335 

336 if isinstance(orelse_var, (bool, int, float, str, np.ndarray)): 

337 orelse_var = tensor_conversion.convert_to_tensor_v2(orelse_var) 

338 

339 if (not tensor_util.is_tf_type(body_var) or 

340 not tensor_util.is_tf_type(orelse_var)): 

341 return 

342 

343 # TODO(mdan): Properly account for CompositeTensors. 

344 if (not hasattr(body_var, 'dtype') or 

345 not hasattr(orelse_var, 'dtype')): 

346 return 

347 

348 if body_var.dtype != orelse_var.dtype: 

349 raise TypeError( 

350 "'{}' has dtype {} in the main branch, but dtype {} in the else" 

351 ' branch'.format(name, body_var.dtype.name, 

352 orelse_var.dtype.name)) 

353 

354 

355def _verify_tf_cond_branch_vars(vars_, symbol_names, branch_name): 

356 """Verifies variables output by a conditional branch for consistency.""" 

357 for name, var_ in zip(symbol_names, vars_): 

358 if isinstance(var_, variables.Undefined): 

359 raise ValueError( 

360 "'{}' must also be initialized in the {} branch".format( 

361 name, branch_name)) 

362 if isinstance(var_, variables.UndefinedReturnValue): 

363 raise ValueError( 

364 'the {} branch must also have a return statement.'.format( 

365 branch_name)) 

366 

367 

368def _verify_tf_cond_vars(body_vars, orelse_vars, symbol_names): 

369 """Verifies variables manipulated by a conditional for consistency.""" 

370 named_vars = zip(symbol_names, body_vars, orelse_vars) 

371 

372 for name, body_var, orelse_var in named_vars: 

373 try: 

374 nest.assert_same_structure(body_var, orelse_var, expand_composites=True) 

375 except (ValueError, TypeError): 

376 # One branch of cond could be a `Tensor`, while the other branch could be 

377 # a `ResourceVariable`. Convert `ResourceVariable`s to `Tensor`s so 

378 # assert_same_structure won't fail. 

379 try: 

380 body_var_tensors = variable_utils.convert_variables_to_tensors(body_var) 

381 orelse_var_tensors = variable_utils.convert_variables_to_tensors( 

382 orelse_var) 

383 nest.assert_same_structure(body_var_tensors, orelse_var_tensors, 

384 expand_composites=True) 

385 except (ValueError, TypeError) as e: 

386 raise TypeError( 

387 "'{}' must have the same nested structure in the main and else" 

388 ' branches:\n\n{}'.format(name, str(e))) from e 

389 nest.map_structure( 

390 functools.partial(verify_single_cond_var, name), body_var, orelse_var) 

391 

392 

393def for_stmt(iter_, extra_test, body, get_state, set_state, symbol_names, opts): 

394 """Functional form of a for statement. 

395 

396 The loop operates on a state, which includes all symbols that are 

397 variant across loop iterations, excluding the variables local to the loop. 

398 

399 For example, given the loop below that calculates the geometric and 

400 arithmetic means or some numbers: 

401 

402 ``` 

403 geo_mean = 1 

404 arith_mean = 0 

405 for i in range(n): 

406 a = numbers[i] 

407 geo_mean *= a 

408 arith_mean += a 

409 ``` 

410 

411 The state is represented by the variables named geo_mean and arith_mean. The 

412 `extra_test`, `body`, `get_state` and `set_state` functions must bind to the 

413 original `geo_mean` and `arith_mean` symbols, using `nonlocal`. 

414 

415 The inputs and outputs of the callables representing the loop blocks are not 

416 explicit - instead, these functions must use nonlocal/global for side effects. 

417 The inputs and outputs are instead controlled by the set_state/get_state 

418 functions. 

419 

420 Args: 

421 iter_: The entity being iterated over. 

422 extra_test: Callable with boolean return type. An additional loop condition. 

423 body: Callable representing the actual loop body. 

424 get_state: Additional callable which can capture additional state (such as 

425 the values of composite symbols). This is only useful when staging the 

426 loop. 

427 set_state: Additional callable which save values captured by get_state back 

428 into the Python environment. This is only useful when staging the loop. 

429 symbol_names: Tuple containing names of the loop variables returned by 

430 get_state. 

431 opts: Optional dict of extra loop parameters. 

432 """ 

433 

434 try: 

435 for_fn = for_loop_registry.lookup(iter_) 

436 except LookupError: 

437 for_fn = _py_for_stmt 

438 

439 # TODO(bwieder): Refactor isinstance(iter_, ragged_tensor.RaggedTensor) to use 

440 # the registry once python/autograph/utils does not depend on dataset_ops. 

441 if tensor_util.is_tf_type(iter_): 

442 if tensors.is_range_tensor(iter_): 

443 for_fn = _tf_range_for_stmt 

444 elif isinstance(iter_, ragged_tensor.RaggedTensor): 

445 for_fn = _tf_ragged_for_stmt 

446 else: 

447 for_fn = _known_len_tf_for_stmt 

448 elif isinstance(iter_, distribute.Iterator): 

449 for_fn = _tf_iterator_for_stmt 

450 elif isinstance(iter_, distribute.Iterable): 

451 # TODO(b/162250181): Use _tf_iterator_for_stmt(iter(iter_)... 

452 for_fn = _tf_distributed_iterable_for_stmt 

453 

454 for_fn(iter_, extra_test, body, get_state, set_state, symbol_names, opts) 

455 

456 

457def _py_for_stmt( 

458 iter_, extra_test, body, get_state, set_state, symbol_names, opts 

459): 

460 """Overload of for_stmt that executes a Python for loop.""" 

461 del get_state, set_state, symbol_names, opts 

462 

463 if __debug__: 

464 checker = _PythonLoopChecker() 

465 before_iteration = checker.before_iteration 

466 after_iteration = checker.after_iteration 

467 before_iteration() 

468 

469 original_body = body 

470 def protected_body(protected_iter): 

471 original_body(protected_iter) 

472 after_iteration() 

473 before_iteration() 

474 body = protected_body 

475 

476 if extra_test is not None: 

477 def guarded_extra_test(): 

478 extra_test_result = extra_test() 

479 try: 

480 # Note: Using try/except and not tensor_util.is_tf_type to avoid 

481 # performance degradation. 

482 return bool(extra_test_result) 

483 except errors_impl.OperatorNotAllowedInGraphError as e: 

484 ag_logging.log( 

485 1, 

486 'Caught error while evaluating loop stop condition', 

487 exc_info=True) 

488 # TODO(mdan): We can pass the location of extra_test and show it here. 

489 raise NotImplementedError( 

490 'break and return statements which depend on a TF condition are not' 

491 ' supported in Python for loops. Did you intend to make it a TF' 

492 ' loop?\nSee ' 

493 'https://github.com/tensorflow/tensorflow/blob/master/tensorflow/' 

494 'python/autograph/g3doc/reference/limitations.md' 

495 '#consistency-of-control-flow-types for more info.') from e 

496 

497 if guarded_extra_test(): 

498 for target in iter_: 

499 body(target) 

500 if not guarded_extra_test(): 

501 break 

502 

503 else: 

504 for target in iter_: 

505 body(target) 

506 

507 

508def _add_max_iterations_hint(opts, n): 

509 # TODO(b/159186914): Remove the safeguard, and always set maximum_iterations. 

510 if control_flow_util.GraphOrParentsInXlaContext(ops.get_default_graph()): 

511 opts['maximum_iterations'] = n 

512 

513 

514def _known_len_tf_for_stmt( 

515 iter_, extra_test, body, get_state, set_state, symbol_names, opts): 

516 """Overload of for_stmt that iterates over TF entities that admit a length.""" 

517 n = py_builtins.len_(iter_) 

518 

519 # TODO(b/117628877): Revisit performance once XLA has the necessary support. 

520 # Note: using a TensorArray creates an extra copy, but can calculate 

521 # gradients more efficiently than StridedSlice. 

522 ta = tensor_array_ops.TensorArray(iter_.dtype, size=n) 

523 iter_ = ta.unstack(iter_) 

524 

525 iterate_index = 0 

526 

527 def aug_get_state(): 

528 return (iterate_index,) + get_state() 

529 

530 def aug_set_state(aug_loop_vars): 

531 nonlocal iterate_index 

532 # TODO(b/171479293): Drop the lint override. 

533 iterate_index, *loop_vars = aug_loop_vars # pylint:disable=unused-variable 

534 # The iteration index is not "output" by the for loop. If the iteration index 

535 # is used outside the loop, it will appear in the loop vars separately. 

536 set_state(loop_vars) 

537 

538 def aug_body(): 

539 nonlocal iterate_index 

540 body(iter_.read(iterate_index)) 

541 iterate_index += 1 

542 

543 def aug_test(): 

544 main_test = iterate_index < n 

545 if extra_test is not None: 

546 return tf_cond.cond(main_test, extra_test, lambda: False) 

547 return main_test 

548 

549 _add_max_iterations_hint(opts, n) 

550 

551 _tf_while_stmt( 

552 aug_test, 

553 aug_body, 

554 aug_get_state, 

555 aug_set_state, 

556 ('<internal iterate>',) + symbol_names, 

557 opts, 

558 ) 

559 

560 

561def _tf_ragged_for_stmt( 

562 iter_, extra_test, body, get_state, set_state, symbol_names, opts): 

563 """Overload of for_stmt that iterates over TF ragged tensors.""" 

564 init_vars = get_state() 

565 verify_loop_init_vars(init_vars, symbol_names) 

566 

567 # TODO(mdan): Move this into len()? Requires eager support. 

568 if iter_.shape and iter_.shape[0] is not None: 

569 n = iter_.shape[0] 

570 else: 

571 n = iter_.row_lengths()[0] 

572 

573 iterate_index = 0 

574 

575 def aug_get_state(): 

576 return (iterate_index,) + get_state() 

577 

578 def aug_set_state(aug_loop_vars): 

579 nonlocal iterate_index 

580 # TODO(b/171479293): Drop the lint override. 

581 iterate_index, *loop_vars = aug_loop_vars # pylint:disable=unused-variable 

582 # The iteration index is not "output" by the for loop. If the iteration index 

583 # is used outside the loop, it will appear in the loop vars separately. 

584 set_state(loop_vars) 

585 

586 def aug_body(): 

587 nonlocal iterate_index 

588 body(iter_[iterate_index]) 

589 iterate_index += 1 

590 

591 def aug_test(): 

592 main_test = iterate_index < n 

593 if extra_test is not None: 

594 return tf_cond.cond(main_test, extra_test, lambda: False) 

595 return main_test 

596 

597 _add_max_iterations_hint(opts, n) 

598 

599 _tf_while_stmt( 

600 aug_test, 

601 aug_body, 

602 aug_get_state, 

603 aug_set_state, 

604 ('<internal iterate>',) + symbol_names, 

605 opts) 

606 

607 

608def _tf_range_for_stmt( 

609 iter_, extra_test, body, get_state, set_state, symbol_names, opts): 

610 """Overload of for_stmt that iterates over a TF range (and elides it).""" 

611 start, limit, delta = iter_.op.inputs 

612 

613 iterate = start 

614 

615 def _value_or(name, var, default): 

616 if (name == opts['iterate_names'] and isinstance(var, variables.Undefined)): 

617 return default 

618 return var 

619 

620 def aug_get_state(): 

621 state_vars = get_state() 

622 state_vars = tuple( 

623 _value_or(name, var, iterate) 

624 for name, var in zip(symbol_names, state_vars)) 

625 return (iterate,) + state_vars 

626 

627 def aug_set_state(aug_loop_vars): 

628 nonlocal iterate 

629 # TODO(b/171479293): Drop the lint override. 

630 iterate, *loop_vars = aug_loop_vars # pylint:disable=unused-variable 

631 # The iteration index is not "output" by the for loop. If the iterate 

632 # is used outside the loop, it will appear in the loop vars separately. 

633 set_state(loop_vars) 

634 

635 def aug_body(): 

636 nonlocal iterate 

637 body(iterate) 

638 iterate += delta 

639 

640 def aug_test(): 

641 # TODO(b/159713842): Remove once constant folding works. 

642 const_delta = tensor_util.constant_value(delta) 

643 if const_delta is not None: 

644 if const_delta >= 0: 

645 main_test = iterate < limit 

646 else: 

647 main_test = iterate > limit 

648 else: 

649 main_test = math_ops.logical_or( 

650 math_ops.logical_and(delta >= 0, iterate < limit), 

651 math_ops.logical_and(delta < 0, iterate > limit)) 

652 

653 if extra_test is not None: 

654 main_test = tf_cond.cond(main_test, extra_test, lambda: False) 

655 return main_test 

656 

657 _add_max_iterations_hint( 

658 opts, 

659 math_ops.cast(misc.get_range_len(start, limit, delta), dtypes.int32)) 

660 

661 _tf_while_stmt( 

662 aug_test, 

663 aug_body, 

664 aug_get_state, 

665 aug_set_state, 

666 ('<internal iterate>',) + symbol_names, 

667 opts) 

668 

669 

670def _tf_iterator_for_stmt( 

671 iter_, extra_test, body, get_state, set_state, symbol_names, opts): 

672 """Overload of for_stmt that iterates over TF Iterators. See for_loop.""" 

673 symbol_names = ('<internal has_next>',) + symbol_names 

674 has_next = True 

675 

676 def aug_get_state(): 

677 return (has_next,) + get_state() 

678 

679 def aug_set_state(aug_loop_vars): 

680 nonlocal has_next 

681 # TODO(b/171479293): Drop the lint override. 

682 has_next, *loop_vars = aug_loop_vars # pylint:disable=unused-variable 

683 set_state(loop_vars) 

684 

685 init_vars = aug_get_state() 

686 verify_loop_init_vars(init_vars, symbol_names) 

687 

688 def aug_body(): 

689 """Main body passed to _tf_while_stmt.""" 

690 nonlocal has_next 

691 opt_iterate = iter_.get_next_as_optional() 

692 has_next = opt_iterate.has_value() 

693 loop_vars = aug_get_state() # updated by set_state() in _tf_while_loop. 

694 

695 def main_path(): 

696 body(opt_iterate.get_value()) 

697 new_loop_vars = aug_get_state() 

698 # Note: this verification duplicates the one performed in tf_while_stmt, 

699 # but needs to be done earlier to prevent the tf.cond from blowing up 

700 # first. 

701 verify_tf_loop_vars( 

702 init_vars, loop_vars, new_loop_vars, symbol_names, opts) 

703 return new_loop_vars 

704 

705 def noop_path(): 

706 return loop_vars 

707 

708 # TODO(mdan): If tf.while_loop supported Optional, this could be avoided. 

709 # Calling set_state so that get_state() _tf_while_loop sees the conditional 

710 # tensors. 

711 aug_set_state( 

712 tf_cond.cond(has_next, main_path, noop_path)) 

713 

714 def aug_test(): 

715 # This value takes a complicated path to get here: 

716 # prev_iteration_body -> get_state -> tf.while_loop (as loop var) 

717 # -> current_iteration_body -> set_state -> has_next 

718 main_test = has_next 

719 if extra_test is not None: 

720 return tf_cond.cond(main_test, extra_test, lambda: False) 

721 return main_test 

722 

723 _tf_while_stmt( 

724 aug_test, 

725 aug_body, 

726 aug_get_state, 

727 aug_set_state, 

728 symbol_names, 

729 opts) 

730 

731 

732def _tf_distributed_iterable_for_stmt( 

733 iter_, extra_test, body, get_state, set_state, symbol_names, opts): 

734 """Overload of for_stmt that iterates over TF distributed datasets.""" 

735 

736 if extra_test is not None: 

737 raise NotImplementedError( 

738 'break and return statements are not yet supported in ' 

739 'for ... in distributed input loops.') 

740 

741 init_vars = get_state() 

742 verify_loop_init_vars(init_vars, symbol_names) 

743 

744 if 'shape_invariants' in opts: 

745 opts['shape_invariants'] = _shape_invariants_mapping_to_positional_list( 

746 opts['shape_invariants'], init_vars) 

747 

748 def reduce_body(loop_vars, iterate): 

749 set_state(loop_vars) 

750 body(iterate) 

751 new_loop_vars = get_state() 

752 verify_tf_loop_vars( 

753 init_vars, loop_vars, new_loop_vars, symbol_names, opts) 

754 return new_loop_vars 

755 

756 set_state(iter_.reduce(init_vars, reduce_body)) 

757 

758 

759def while_stmt(test, body, get_state, set_state, symbol_names, opts): 

760 """Functional form of a while statement. 

761 

762 The loop operates on a so-called state, which includes all symbols that are 

763 variant across loop iterations. In what follows we refer to state as either 

764 a tuple of entities that represent an actual state, or a list of arguments 

765 of the corresponding types. 

766 

767 The inputs and outputs of the callables representing the loop blocks are not 

768 explicit - instead, these functions must use nonlocal/global for side effects. 

769 The inputs and outputs are instead controlled by the set_state/get_state 

770 functions. 

771 

772 Args: 

773 test: Callable with boolean return type. The loop condition. 

774 body: Callable representing the actual loop body. 

775 get_state: Additional callable which can capture additional state (such as 

776 the values of composite symbols). This is only useful when staging the 

777 loop. 

778 set_state: Additional callable which save values captured by get_state back 

779 into the Python environment. This is only useful when staging the loop. 

780 symbol_names: Tuple containing the names of all loop variables. 

781 opts: Optional dict of extra loop parameters. 

782 

783 Returns: 

784 Tuple containing the final state. 

785 """ 

786 

787 # Evaluate the initial test once in order to do the dispatch. The evaluation 

788 # is isolated to minimize unwanted side effects. 

789 # TODO(mdan): Do a full iteration - some state types might lower to Tensor. 

790 with func_graph.FuncGraph('tmp').as_default(): 

791 init_test = test() 

792 

793 # TensorFlow: Multiple evaluations are acceptable in this case, so we're fine 

794 # with the re-evaluation of `test` that `_tf_while_stmt` will make. 

795 if tensors.is_dense_tensor(init_test): 

796 _tf_while_stmt(test, body, get_state, set_state, symbol_names, opts) 

797 return 

798 

799 # Normal Python: We already consumed one evaluation of `test`; consistently, 

800 # unroll one iteration before dispatching to a normal loop. 

801 # TODO(mdan): Push the "init_test" value via opts into _py_while_stmt? 

802 if not init_test: 

803 return 

804 body() 

805 

806 _py_while_stmt(test, body, get_state, set_state, opts) 

807 

808 

809class _PythonLoopChecker(object): 

810 """Verifies Python loops for TF-specific limits.""" 

811 

812 __slots__ = ( 

813 'iterations', 

814 'check_inefficient_unroll', 

815 'check_op_count_after_iteration', 

816 'ops_before_iteration', 

817 ) 

818 

819 def __init__(self): 

820 self.iterations = 1 

821 self.check_inefficient_unroll = WARN_INEFFICIENT_UNROLL 

822 

823 # Triggered when we decided to test the op counts. 

824 self.check_op_count_after_iteration = False 

825 

826 def _get_ops(self): 

827 return set(ops.get_default_graph().get_operations()) 

828 

829 def _check_unroll_limits(self): 

830 if self.iterations > PYTHON_MAX_ITERATIONS: 

831 raise ValueError('iteration limit exceeded') 

832 

833 def _stop_checking_inefficient_unroll(self): 

834 self.check_inefficient_unroll = False 

835 self.check_op_count_after_iteration = False 

836 self.ops_before_iteration = None 

837 

838 def _verify_inefficient_unroll(self): 

839 """Checks for possibly-inefficient creation of ops in a Python loop.""" 

840 assert self.ops_before_iteration is not None 

841 ops_after_iteration = self._get_ops() 

842 new_ops = tuple( 

843 op for op in ops_after_iteration if op not in self.ops_before_iteration) 

844 

845 if len(new_ops) < INEFFICIENT_UNROLL_MIN_OPS: 

846 return False 

847 

848 ag_logging.warning( 

849 'Large unrolled loop detected. Did you mean to use a TF loop?' 

850 ' The following ops were created after iteration %s: %s' 

851 '\nSee' 

852 ' https://github.com/tensorflow/tensorflow/blob/master/' 

853 'tensorflow/python/autograph/g3doc/reference/common_errors.md' 

854 '#warning-large-unrolled-loop-detected' 

855 '\n' 

856 'Location:' 

857 '\n%s' 

858 '', self.iterations, new_ops, '\n'.join(traceback.format_stack())) 

859 return True 

860 

861 def before_iteration(self): 

862 """Called before each iteration in a Python loop.""" 

863 if (self.check_inefficient_unroll and 

864 self.iterations > INEFFICIENT_UNROLL_MIN_ITERATIONS): 

865 self.ops_before_iteration = self._get_ops() 

866 self.check_op_count_after_iteration = True 

867 

868 def after_iteration(self): 

869 """Called after each iteration in a Python loop.""" 

870 self.iterations += 1 

871 

872 self._check_unroll_limits() 

873 

874 if self.check_op_count_after_iteration: 

875 did_warn = self._verify_inefficient_unroll() 

876 if did_warn: 

877 self._stop_checking_inefficient_unroll() # Only warn once. 

878 elif self.iterations > INEFFICIENT_UNROLL_MIN_ITERATIONS + 3: 

879 # Once deciding to check the op counts, only do it for a few iterations. 

880 self._stop_checking_inefficient_unroll() 

881 

882 

883def _py_while_stmt(test, body, get_state, set_state, opts): 

884 """Overload of while_stmt that executes a Python while loop.""" 

885 del opts, get_state, set_state 

886 

887 if __debug__: 

888 checker = _PythonLoopChecker() 

889 before_iteration = checker.before_iteration 

890 after_iteration = checker.after_iteration 

891 before_iteration() 

892 

893 original_body = body 

894 def protected_body(): 

895 original_body() 

896 after_iteration() 

897 before_iteration() 

898 body = protected_body 

899 

900 def guarded_test(): 

901 test_result = test() 

902 try: 

903 # Note: Using try/except and not tensor_util.is_tf_type to avoid 

904 # performance degradation. 

905 return bool(test_result) 

906 except errors_impl.OperatorNotAllowedInGraphError as e: 

907 ag_logging.log( 

908 1, 

909 'Caught error while evaluating while loop condition', 

910 exc_info=True) 

911 # TODO(mdan): distinguish beteen these two cases.