# Path: deep_verify/src/common.py
# def with_explicit_update(volatile_value):
#   def materialise(value):
# def targeted_objective(final_w, final_b, labels):
# def concave_max_binsearch(fn, lb, ub, num_iter=20):
# def _prod(lst):
# def convolution(x, kernel, padding, strides):
# def conv_transpose(y, kernel, result_shape, padding, strides):
# def avgpool_transpose(y, result_shape, kernel_shape, strides):
# def conv_broadcast(x, kernel_shape, padding, strides):
# def conv2d_broadcast(x, kernel_height, kernel_width, padding, strides):
# def conv1d_broadcast(x, kernel_length, padding, stride):
# def conv_reduce_sum(x, result_shape, padding, strides):
# def conv2d_reduce_sum(x, input_height, input_width, padding, strides):
# def conv1d_reduce_sum(x, input_length, padding, stride):
# 
# Path: deep_verify/src/formulations/standard/standard_layer_calcs.py
# def _reduce_softmax(v, axis, inverse_temperature=None):
# def _conj(nl, mu, lam, nominal, lb, ub, parameters=None,
#           inverse_temperature=None):
#     def siginv(y):
# def max_linear(coefficients, lb, ub, axis=None, keepdims=False,
#                inverse_temperature=None):
# def linear_dual_objective(lam_in, activation_coeffs, dual_obj_bias, lb, ub,
#                           inverse_temperature=None):
# def _prod(lst):
# def maxpool_layer_dual_objective(kernel_shape, strides, with_relu,
#                                  mu_in, lam_out, lb, ub, nominal=None):
#   def fn(y):
# def activation_layer_dual_objective(nl, mu_in, lam_out, lb, ub, nominal=None,
#                                     parameters=None, inverse_temperature=None):


from absl.testing import parameterized
from deep_verify.src import common
from deep_verify.src.formulations.standard import standard_layer_calcs
from interval_bound_propagation import layer_utils
import numpy as np
import sonnet as snt
import tensorflow as tf
    lb_max = tf.expand_dims(tf.expand_dims(lb_max, 1), 2)

    def dual_obj_given_max_at_yi(c_k):
      # Evaluate max (mu^T z - lam y_i) at y_i = c_k.
      bc = tf.zeros_like(mu_i + c_k)  # tf.where doesn't broadcast
      dual_obj = standard_layer_calcs.max_linear(
          mu_i + bc, lb_i + bc,
          tf.minimum(c_k, ub_i), axis=[2, 3], keepdims=True)
      dual_obj -= tf.maximum(lb_i * mu_i, tf.minimum(c_k, ub_i) * mu_i)
      dual_obj += (mu_i - lam) * c_k
      # Only consider this y if it's in the permitted range for a maximal y_i.
      feasible = tf.logical_and(lb_max <= c_k, c_k <= ub_i + bc)
      dual_obj = tf.where(feasible, dual_obj, -1.e8 + bc)
      # Take maximum over all i, k.
      return tf.reduce_max(dual_obj, axis=[2, 3, 4, 5])

    # Evaluate max (mu^T z - lam y_i) at y_i = max_j lb_j.
    dual_obj_lb_max = dual_obj_given_max_at_yi(lb_max)
    # Evaluate max (mu^T z - lam y_i) at y_i = ub_k.
    dual_obj_ub_k = dual_obj_given_max_at_yi(ub_k)

    dual_obj_max = tf.maximum(dual_obj_lb_max, dual_obj_ub_k)

    if with_relu:
      # Also consider the case in which all y_i are <= 0,
      # so relu_maxpool is zero.
      # This leads us to maximise (over z_j in [lb_j, min{0, ub_j}])::
      #   mu^T z  -  lam 0

      bc = tf.zeros_like(mu_i)  # tf.where doesn't broadcast

# Path: deep_verify/src/common.py
# def with_explicit_update(volatile_value):
#   def materialise(value):
# def targeted_objective(final_w, final_b, labels):
# def concave_max_binsearch(fn, lb, ub, num_iter=20):
# def _prod(lst):
# def convolution(x, kernel, padding, strides):
# def conv_transpose(y, kernel, result_shape, padding, strides):
# def avgpool_transpose(y, result_shape, kernel_shape, strides):
# def conv_broadcast(x, kernel_shape, padding, strides):
# def conv2d_broadcast(x, kernel_height, kernel_width, padding, strides):
# def conv1d_broadcast(x, kernel_length, padding, stride):
# def conv_reduce_sum(x, result_shape, padding, strides):
# def conv2d_reduce_sum(x, input_height, input_width, padding, strides):
# def conv1d_reduce_sum(x, input_length, padding, stride):
# 
# Path: deep_verify/src/formulations/standard/standard_layer_calcs.py
# def _reduce_softmax(v, axis, inverse_temperature=None):
# def _conj(nl, mu, lam, nominal, lb, ub, parameters=None,
#           inverse_temperature=None):
#     def siginv(y):
# def max_linear(coefficients, lb, ub, axis=None, keepdims=False,
#                inverse_temperature=None):
# def linear_dual_objective(lam_in, activation_coeffs, dual_obj_bias, lb, ub,
#                           inverse_temperature=None):
# def _prod(lst):
# def maxpool_layer_dual_objective(kernel_shape, strides, with_relu,
#                                  mu_in, lam_out, lb, ub, nominal=None):
#   def fn(y):
# def activation_layer_dual_objective(nl, mu_in, lam_out, lb, ub, nominal=None,
#                                     parameters=None, inverse_temperature=None):


# Path: deep_verify/tests/formulations/standard/standard_layer_calcs_test.py
from absl.testing import parameterized
from deep_verify.src import common
from deep_verify.src.formulations.standard import standard_layer_calcs
from interval_bound_propagation import layer_utils
import numpy as np
import sonnet as snt
import tensorflow as tf

    lb_max = tf.expand_dims(tf.expand_dims(lb_max, 1), 2)

    def dual_obj_given_max_at_yi(c_k):
      # Evaluate max (mu^T z - lam y_i) at y_i = c_k.
      bc = tf.zeros_like(mu_i + c_k)  # tf.where doesn't broadcast
      dual_obj = standard_layer_calcs.max_linear(
          mu_i + bc, lb_i + bc,
          tf.minimum(c_k, ub_i), axis=[2, 3], keepdims=True)
      dual_obj -= tf.maximum(lb_i * mu_i, tf.minimum(c_k, ub_i) * mu_i)
      dual_obj += (mu_i - lam) * c_k
      # Only consider this y if it's in the permitted range for a maximal y_i.
      feasible = tf.logical_and(lb_max <= c_k, c_k <= ub_i + bc)
      dual_obj = tf.where(feasible, dual_obj, -1.e8 + bc)
      # Take maximum over all i, k.
      return tf.reduce_max(dual_obj, axis=[2, 3, 4, 5])

    # Evaluate max (mu^T z - lam y_i) at y_i = max_j lb_j.
    dual_obj_lb_max = dual_obj_given_max_at_yi(lb_max)
    # Evaluate max (mu^T z - lam y_i) at y_i = ub_k.
    dual_obj_ub_k = dual_obj_given_max_at_yi(ub_k)

    dual_obj_max = tf.maximum(dual_obj_lb_max, dual_obj_ub_k)

    if with_relu:
      # Also consider the case in which all y_i are <= 0,
      # so relu_maxpool is zero.
      # This leads us to maximise (over z_j in [lb_j, min{0, ub_j}])::
      #   mu^T z  -  lam 0

      bc = tf.zeros_like(mu_i)  # tf.where doesn't broadcast

      dual_obj_zero = standard_layer_calcs.max_linear(