# 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):


import functools
import operator
import tensorflow as tf
from deep_verify.src import common
    """Optimal dual objective, conditional on the value of the maxpool.

    For each input (kernel element) i, for the given y_i,
    maximises (over z_j in [lb_j, min{y_i, ub_j}] and constraining z_i=y_i)::
      mu^T z  -  lam y_i

    This will be infeasible if y_i < lb_j for some j, (also if y_i < 0 in the
    case of relu+maxpool), so maxpool cannot be attained at input i. The
    returned tensor is unspecified for such elements.

    Args:
      y: (N+4)D tensor of shape (num_classes, batch_size,
        output_height, output_width,
        kernel_height * kernel_width, layer_channels) containing, for each
        input (kernel element) i, a value of maxpool assumed to be attained
        at input i.

    Returns:
      Tensor of same shape as `y` containing, for each input (kernel element) i,
        the optimal value of the dual objective, conditional the maxpool being
        equal to `y` with the maximum attained at input i.
    """
    y_i = tf.expand_dims(y, nd+3)
    # Maximise sum_{j!=i} mu_j y_j  where y_j <= y_i for all j!=i.
    obj = max_linear(mask * mu_j, lb_j, tf.minimum(ub_j, y_i), axis=(nd+3))
    return obj + (mu_bc - lam_bc) * y

  lb_max = tf.reduce_max(lb_bc, axis=(nd+2), keepdims=True)
  if with_relu:
    lb_max = tf.maximum(lb_max, 0.)

# 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
import functools
import operator
import tensorflow as tf
from deep_verify.src import common

    """Optimal dual objective, conditional on the value of the maxpool.

    For each input (kernel element) i, for the given y_i,
    maximises (over z_j in [lb_j, min{y_i, ub_j}] and constraining z_i=y_i)::
      mu^T z  -  lam y_i

    This will be infeasible if y_i < lb_j for some j, (also if y_i < 0 in the
    case of relu+maxpool), so maxpool cannot be attained at input i. The
    returned tensor is unspecified for such elements.

    Args:
      y: (N+4)D tensor of shape (num_classes, batch_size,
        output_height, output_width,
        kernel_height * kernel_width, layer_channels) containing, for each
        input (kernel element) i, a value of maxpool assumed to be attained
        at input i.

    Returns:
      Tensor of same shape as `y` containing, for each input (kernel element) i,
        the optimal value of the dual objective, conditional the maxpool being
        equal to `y` with the maximum attained at input i.
    """
    y_i = tf.expand_dims(y, nd+3)
    # Maximise sum_{j!=i} mu_j y_j  where y_j <= y_i for all j!=i.
    obj = max_linear(mask * mu_j, lb_j, tf.minimum(ub_j, y_i), axis=(nd+3))
    return obj + (mu_bc - lam_bc) * y

  lb_max = tf.reduce_max(lb_bc, axis=(nd+2), keepdims=True)
  if with_relu:
    lb_max = tf.maximum(lb_max, 0.)

  _, attained = common.concave_max_binsearch(