text: # Copyright 2018 The TensorFlow Constrained Optimization 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.
  # ==============================================================================
  """Tests for operations.py."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow.compat.v2 as tf
  
  from tensorflow_constrained_optimization.python import graph_and_eager_test_case
  from tensorflow_constrained_optimization.python.rates import defaults
  from tensorflow_constrained_optimization.python.rates import deferred_tensor
  from tensorflow_constrained_optimization.python.rates import operations
  # Placeholder for internal import.
  
  
  # @run_all_tests_in_graph_and_eager_modes
  class OperationsTest(graph_and_eager_test_case.GraphAndEagerTestCase):
    """Tests for `Expression`-manipulation functions."""
  
    def _evaluate_expression(self, expression, extra_update_ops_fn=None):
      """Evaluates and returns both portions of an Expression.
  
      Args:
        expression: `Expression` to evaluate.
        extra_update_ops_fn: function that takes an `EvaluationMemoizer` and the
          list of `DeferredVariables`, and returns a list of ops to execute before
          evaluation.
  
      Returns:
        A pair (penalty,constraint) containing the values of the penalty and
        constraint portions of the `Expression`.
      """
      structure_memoizer = {
          defaults.DENOMINATOR_LOWER_BOUND_KEY: 0.0,
          defaults.GLOBAL_STEP_KEY: tf.Variable(0, dtype=tf.int32),
          defaults.VARIABLE_FN_KEY: tf.Variable
      }
  
      penalty_value = expression.penalty_expression.evaluate(structure_memoizer)
      constraint_value = expression.constraint_expression.evaluate(
          structure_memoizer)
  
      # We need to explicitly create the variables before creating the wrapped
      # session.
      variables = deferred_tensor.DeferredVariableList(
          penalty_value.variables + constraint_value.variables).list
      for variable in variables:
        variable.create(structure_memoizer)
  
      def update_ops_fn():
        if not extra_update_ops_fn:
          update_ops = []
        else:
          update_ops = extra_update_ops_fn(structure_memoizer, variables)
        for variable in variables:
          update_ops += variable.update_ops(structure_memoizer)
        return update_ops
  
      with self.wrapped_session() as session:
        session.run_ops(update_ops_fn)
        return [
            session.run(penalty_value(structure_memoizer)),
            session.run(constraint_value(structure_memoizer))
        ]
  
    def test_wrap_rate_raises(self):
      """Make sure that wrap_rate() raises if given an `Expression`."""
      penalty_value = 3.1
      constraint_value = 2.7
  
      wrapped = operations.wrap_rate(penalty_value, constraint_value)
      # We should raise a TypeError if either or both of the parameters to
      # wrap_rate() are Expressions.
      with self.assertRaises(TypeError):
        operations.wrap_rate(penalty_value, wrapped)
      with self.assertRaises(TypeError):
        operations.wrap_rate(wrapped, constraint_value)
      with self.assertRaises(TypeError):
        operations.wrap_rate(wrapped, wrapped)
  
    def test_wrap_rate_penalty(self):
      """Make sure that wrap_rate() works correctly when given one parameter."""
      penalty_value = 3.1
      penalty_tensor = tf.constant(penalty_value, dtype=tf.float32)
  
      wrapped = operations.wrap_rate(penalty_tensor)
      actual_penalty_value, actual_constraint_value = self._evaluate_expression(
          wrapped)
      # Both the penalty and the constraint portions of the Expression were
      # initialized to penalty_tensor.
      self.assertNear(penalty_value, actual_penalty_value, err=1e-6)
      self.assertNear(penalty_value, actual_constraint_value, err=1e-6)
  
    def test_wrap_rate_penalty_and_constraint(self):
      """Make sure that wrap_rate() works correctly when given two parameters."""
      penalty_value = 3.1
      constraint_value = 2.7
      penalty_tensor = tf.constant(penalty_value, dtype=tf.float32)
      constraint_tensor = tf.constant(constraint_value, dtype=tf.float32)
  
      wrapped = operations.wrap_rate(penalty_tensor, constraint_tensor)
      wrapped = operations.wrap_rate(penalty_tensor)
      actual_penalty_value, actual_constraint_value = self._evaluate_expression(
          wrapped)
      # The penalty and constraint portions of the Expression were initialized to
      # penalty_tensor and constraint_tensor, respectively.
      self.assertNear(penalty_value, actual_penalty_value, err=1e-6)
      self.assertNear(penalty_value, actual_constraint_value, err=1e-6)
  
    def test_upper_bound_raises_on_empty_list(self):
      """Make sure that upper_bound() raises for an empty expressions list."""
      with self.assertRaises(ValueError):
        operations.upper_bound([])
  
    def test_upper_bound_raises_on_tensor(self):
      """Make sure that upper_bound() raises when given a non-Expression."""
      value1 = 3.1
      value2 = 2.7
      tensor1 = tf.constant(value1, dtype=tf.float32)
      expression2 = operations.wrap_rate(tf.constant(value2, dtype=tf.float64))
  
      # List element is a Tensor, instead of an Expression.
      with self.assertRaises(TypeError):
        operations.upper_bound([tensor1, expression2])
  
    def test_upper_bound_raises_on_maximization(self):
      """Make sure that upper_bound() raises if it's maximized."""
      bounded = -operations.upper_bound([operations.wrap_rate(1.0)])
      # All three of "penalty_expression", "constraint_expression" and
      # "extra_constraints" should raise.
      with self.assertRaises(RuntimeError):
        _ = bounded.penalty_expression
      with self.assertRaises(RuntimeError):
        _ = bounded.constraint_expression
      with self.assertRaises(RuntimeError):
        _ = bounded.extra_constraints
  
    def test_upper_bound(self):
      """Make sure that upper_bound() creates the correct Expression."""
      values = [0.8, 3.1, -1.6, 2.7]
      bound_value = 1.4
      expressions = [operations.wrap_rate(value) for value in values]
  
      bounded = operations.upper_bound(expressions)
  
      # Before evaluating any expressions, we'll assign "bound_value" to the slack
      # variable, so that we can make sure that the same slack variable is being
      # used for all of the constraints.
      def update_ops_fn(structure_memoizer, variables):
        upper_bound_tensor = None
        for variable in variables:
          tensor = variable(structure_memoizer)
          if tensor.name.startswith("tfco_upper_bound"):
            self.assertIsNone(upper_bound_tensor)
            upper_bound_tensor = tensor
        self.assertIsNotNone(upper_bound_tensor)
        return [upper_bound_tensor.assign(bound_value)]
  
      # Extract the set of constraints, and make sure that there is one for each
      # quantity that is being bounded.
      self.assertEqual(len(values), len(bounded.extra_constraints))
      constraints = list(bounded.extra_constraints)
  
      # Evaluate the constraint expressions.
      actual_values = []
      for constraint in constraints:
        actual_penalty_value, actual_constraint_value = self._evaluate_expression(
            constraint.expression, update_ops_fn)
        self.assertEqual(actual_penalty_value, actual_constraint_value)
        actual_values.append(actual_penalty_value)
      # Constraints take the form expression <= 0, and these are upper-bound
      # constraints, so we're expecting:
      #     value1 - bound_value <= 0
      #     value2 - bound_value <= 0
      #       ....
      # We sort the constraint expression values since they occur in no particular
      # order.
      actual_values = sorted(actual_values)
      expected_values = sorted(value - bound_value for value in values)
      self.assertAllClose(expected_values, actual_values, rtol=0, atol=1e-6)
  
    def test_lower_bound_raises_on_empty_list(self):
      """Make sure that lower_bound() raises for an empty expressions list."""
      with self.assertRaises(ValueError):
        operations.lower_bound([])
  
    def test_lower_bound_raises_on_tensor(self):
      """Make sure that lower_bound() raises when given a non-Expression."""
      value1 = 3.1
      value2 = 2.7
      tensor1 = tf.constant(value1, dtype=tf.float32)
      expression2 = operations.wrap_rate(tf.constant(value2, dtype=tf.float64))
  
      # List element is a Tensor, instead of an Expression.
      with self.assertRaises(TypeError):
        operations.lower_bound([tensor1, expression2])
  
    def test_lower_bound_raises_on_minimization(self):
      """Make sure that lower_bound() raises if it's minimized."""
      bounded = operations.lower_bound([operations.wrap_rate(1.0)])
      # All three of "penalty_expression", "constraint_expression" and
      # "extra_constraints" should raise.
      with self.assertRaises(RuntimeError):
        _ = bounded.penalty_expression
      with self.assertRaises(RuntimeError):
        _ = bounded.constraint_expression
      with self.assertRaises(RuntimeError):
        _ = bounded.extra_constraints
  
    def test_lower_bound(self):
      """Make sure that lower_bound() creates the correct Expression."""
      values = [0.8, 3.1, -1.6, 2.7]
      bound_value = 1.4
      expressions = [operations.wrap_rate(value) for value in values]
  
      # We need to negate "bounded" since it's implicitly being minimized, and we
      # cannot minimize a lower bound.
      bounded = -operations.lower_bound(expressions)
  
      # Before evaluating any expressions, we'll assign "bound_value" to the slack
      # variable, so that we can make sure that the same slack variable is being
      # used for all of the constraints.
      def update_ops_fn(structure_memoizer, variables):
        lower_bound_tensor = None
        for variable in variables:
          tensor = variable(structure_memoizer)
          if tensor.name.startswith("tfco_lower_bound"):
            self.assertIsNone(lower_bound_tensor)
            lower_bound_tensor = tensor
        self.assertIsNotNone(lower_bound_tensor)
        return [lower_bound_tensor.assign(bound_value)]
  
      # Extract the set of constraints, and make sure that there is one for each
      # quantity that is being bounded.
      self.assertEqual(len(values), len(bounded.extra_constraints))
      constraints = list(bounded.extra_constraints)
  
      # Evaluate the constraint expressions.
      actual_values = []
      for constraint in constraints:
        actual_penalty_value, actual_constraint_value = self._evaluate_expression(
            constraint.expression, update_ops_fn)
        self.assertEqual(actual_penalty_value, actual_constraint_value)
        actual_values.append(actual_penalty_value)
      # Constraints take the form expression <= 0, and these are lower-bound
      # constraints, so we're expecting:
      #     bound_value - value1 <= 0
      #     bound_value - value2 <= 0
      #       ....
      # We sort the constraint expression values since they occur in no particular
      # order.
      actual_values = sorted(actual_values)
      expected_values = sorted(bound_value - value for value in values)
      self.assertAllClose(expected_values, actual_values, rtol=0, atol=1e-6)
  
  
  if __name__ == "__main__":
    tf.test.main()
  