Function _tf_while_stmt

tensorflow/python/autograph/operators/control_flow.py:736–772

Overload of while_stmt that stages a TF while_stmt.

(test, body, get_state, set_state, init_vars,
                   basic_symbol_names, composite_symbol_names, opts)

# composite_symbol_names=None and call _verify_tf_loop_vars(...) itself.  We can
# remove these arguments once all callers do that.
def _tf_while_stmt(test, body, get_state, set_state, init_vars,
                   basic_symbol_names, composite_symbol_names, opts):
  """Overload of while_stmt that stages a TF while_stmt."""
  _disallow_undefs_into_loop(*init_vars)

  if opts is None:
    opts = {}

  # TODO(mdan): Simplify this.
  loop_vars_slice = slice(len(init_vars))
  state_slice = slice(len(init_vars), None)

  def aug_test(*aug_loop_vars):
    state = aug_loop_vars[state_slice]
    set_state(state)
    return test(*aug_loop_vars[loop_vars_slice])

  def aug_body(*aug_loop_vars):
    state = aug_loop_vars[state_slice]
    set_state(state)
    loop_vars = body(*aug_loop_vars[loop_vars_slice])
    new_state = loop_vars + get_state()
    _verify_tf_loop_vars(aug_loop_vars, new_state, basic_symbol_names,
                         composite_symbol_names)

    return new_state

  # Non-v2 while_loop unpacks the results when there is only one return value.
  # This enforces consistency across versions.
  opts['return_same_structure'] = True

  aug_init_vars = init_vars + get_state()
  final_aug_vars = control_flow_ops.while_loop(aug_test, aug_body,
                                               aug_init_vars, **opts)
  final_state = final_aug_vars[state_slice]
  set_state(final_state)
  return final_aug_vars[loop_vars_slice]


class _PythonLoopChecker(object):