text: # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Control flow statements: loops, conditionals, etc."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from tensorflow.contrib.autograph.utils import builtins
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import gen_math_ops
  
  # TODO (mdan): Rename _loop to _stmt to follow Python nomenclature. id:679
  # https://github.com/imdone/tensorflow/issues/680
  # TODO (mdan): Rename arguments to match the AST names. id:528
  # https://github.com/imdone/tensorflow/issues/529
  
  
  def for_loop(iterated, extra_cond, loop_body, init_state):
    """Functional form of a for statement.
  
    The loop operates on a so-called state, which includes all symbols that are
    variant across loop iterations, excluding the iterate. In what follows we
    refer to state as either a tuple of entities that represent an actual state,
    or a list of arguments of the corresponding types.
  
    Args:
      iterated: The entity being iterated over.
      extra_cond: Callable with the state as arguments, and boolean return type.
          An additionnal loop condition.
      loop_body: Callable with the iterate and the state as arguments, and
          state as return type. The actual loop body.
      init_state: Tuple containing the initial state.
  
    Returns:
      Tuple containing the final state.
    """
    if tensor_util.is_tensor(iterated):
      return _known_len_for_loop(iterated, extra_cond, loop_body, init_state)
    elif isinstance(iterated, dataset_ops.Dataset):
      return _dataset_for_loop(iterated, extra_cond, loop_body, init_state)
    else:
      return _py_for_loop(iterated, extra_cond, loop_body, init_state)
  
  
  def _py_for_loop(iterated, extra_cond, loop_body, init_state):
    """Overload of for_loop that executes a Python for loop."""
    state = init_state
    for iterate in iterated:
      if not extra_cond(*state):
        break
      state = loop_body(iterate, *state)
  
    # TODO (mdan): Remove this special case. id:520
    # https://github.com/imdone/tensorflow/issues/521
    if len(state) == 1:
      return state[0]
    return state
  
  
  def _known_len_for_loop(iterated, extra_cond, loop_body, init_state):
    """Overload of for_loop that iterates over objects that define a length."""
    n = builtins.dynamic_len(iterated)
  
    def while_body(iterate_index, *state):
      iterate = iterated[iterate_index]
      new_state = loop_body(iterate, *state)
      return (iterate_index + 1,) + new_state
  
    def while_cond(iterate_index, *state):
      return gen_math_ops.logical_and(iterate_index < n, extra_cond(*state))
  
    results = while_loop(
        while_cond,
        while_body,
        init_state=(0,) + init_state,
        extra_deps=(iterated,),
        opts=dict(maximum_iterations=n))
    # Dropping the iteration index because it's not syntactically visible.
    results = results[1:]
  
    # TODO (mdan): Remove this special case. id:976
    # https://github.com/imdone/tensorflow/issues/977
    if len(results) == 1:
      return results[0]
    return results
  
  
  def _dataset_for_loop(ds, extra_cond, loop_body, init_state):
    """Overload of for_loop that iterates over TF Datasets."""
    # Because Datsets only expose get_next, in the style of Python iterators,
    # we are forced to unpack the loop as:
    #
    # epoch_number, iterate = ds.get_next()
    # while epoch_number < 2:
    #   <body>
    #   epoch_number, iterate = ds.get_next()
    epoch_numbers = dataset_ops.Dataset.range(2)
    def tag_with(ds, tag):
      return dataset_ops.Dataset.zip(
          (dataset_ops.Dataset.from_tensors(tag).repeat(), ds))
    ds_with_epoch = epoch_numbers.flat_map(lambda i: tag_with(ds, i))
  
    iterator = ds_with_epoch.make_initializable_iterator()
    with ops.control_dependencies((iterator.initializer,)):
      epoch_number, iterate = iterator.get_next()
  
      def while_body(epoch_number, iterate, *state):
        new_state = loop_body(iterate, *state)
        epoch_number, iterate = iterator.get_next()
        return (epoch_number, iterate) + new_state
  
      def while_cond(epoch_number, iterate, *state):
        del iterate
        return gen_math_ops.logical_and(epoch_number < 1, extra_cond(*state))
  
      results = while_loop(
          while_cond,
          while_body,
          init_state=(epoch_number, iterate) + init_state,
          extra_deps=())
    # Dropping the epoch number and iterate because they are not not syntactically
    # visible.
    results = results[2:]
  
    # TODO (mdan): Remove this special case. id:630
    # https://github.com/imdone/tensorflow/issues/631
    if len(results) == 1:
      return results[0]
    return results
  
  
  def while_loop(loop_cond, loop_body, init_state, extra_deps, opts=None):
    """Functional form of a while statement.
  
    The loop operates on a so-called state, which includes all symbols that are
    variant across loop iterations. In what follows we refer to state as either
    a tuple of entities that represent an actual state, or a list of arguments
    of the corresponding types.
  
    Args:
      loop_cond: Callable with the state as arguments, and boolean return type.
          The loop condition.
      loop_body: Callable with the state as arguments, and state as return type.
          The actual loop body.
      init_state: Tuple containing the initial state.
      extra_deps: Tuple containing additional entities on which the loop may
          depend, such as loop invariants referenced by loop_cond. Used
          exclusively for dispatch control.
      opts: Optional dict of extra loop parameters.
  
    Returns:
      Tuple containing the final state.
    """
    # TODO (mdan): Consider adding a generic mechanism for dynamic dispatch. id:682
    # https://github.com/imdone/tensorflow/issues/683
    # That could be somethins as simple as a collection of dispatch rules, with
    # some prioritization.
    if any(tensor_util.is_tensor(v) for v in init_state + extra_deps):
      return _tf_while_loop(loop_cond, loop_body, init_state, opts)
    else:
      return _py_while_loop(loop_cond, loop_body, init_state, opts)
  
  
  def _tf_while_loop(loop_cond, loop_body, init_state, opts):
    """Overload of while_loop that stages a TF while_loop."""
    if opts is None:
      opts = {}
    return control_flow_ops.while_loop(loop_cond, loop_body, init_state, **opts)
  
  
  def _py_while_loop(loop_cond, loop_body, init_state, opts):
    """Overload of while_loop that executes a Python while loop."""
    del opts
    state = init_state
    while loop_cond(*state):
      state = loop_body(*state)
    return state
  
  
  def if_stmt(cond, body, orelse):
    """Functional form of an if statement.
  
    Args:
      cond: Boolean.
      body: Callable with no arguments, and outputs of the positive (if) branch
          as return type.
      orelse: Callable with no arguments, and outputs of the negative (else)
          branch as return type.
  
    Returns:
      Tuple containing the statement outputs.
    """
    if tensor_util.is_tensor(cond):
      return _tf_if_stmt(cond, body, orelse)
    else:
      return _py_if_stmt(cond, body, orelse)
  
  
  def _tf_if_stmt(cond, body, orelse):
    """Overload of if_stmt that stages a TF cond."""
    return control_flow_ops.cond(cond, body, orelse)
  
  
  def _py_if_stmt(cond, body, orelse):
    """Overload of if_stmt that executes a Python if statement."""
    return body() if cond else orelse()
  