# Path: interval_bound_propagation/src/bounds.py
# class AbstractBounds(object):
# class IntervalBounds(AbstractBounds):
#   def __init__(self):
#   def convert(cls, bounds):
#   def shape(self):
#   def concretize(self):
#   def _raise_not_implemented(self, name):
#   def apply_linear(self, wrapper, w, b):  # pylint: disable=unused-argument
#   def apply_conv1d(self, wrapper, w, b, padding, stride):  # pylint: disable=unused-argument
#   def apply_conv2d(self, wrapper, w, b, padding, strides):  # pylint: disable=unused-argument
#   def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):  # pylint: disable=unused-argument
#   def apply_piecewise_monotonic_fn(self, wrapper, fn, boundaries, *args):  # pylint: disable=unused-argument
#   def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):  # pylint: disable=unused-argument
#   def apply_batch_reshape(self, wrapper, shape):  # pylint: disable=unused-argument
#   def apply_softmax(self, wrapper):  # pylint: disable=unused-argument
#   def update_cache_op(self):
#   def enable_caching(self):
#   def _set_up_cache(self):
#   def _cache_with_update_op(self, tensor):
#   def __init__(self, lower, upper):
#   def lower(self):
#   def upper(self):
#   def shape(self):
#   def __iter__(self):
#   def convert(cls, bounds):
#   def apply_linear(self, wrapper, w, b):
#   def apply_conv1d(self, wrapper, w, b, padding, stride):
#   def apply_conv2d(self, wrapper, w, b, padding, strides):
#   def _affine(self, w, b, fn, **kwargs):
#   def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):
#   def apply_piecewise_monotonic_fn(self, wrapper, fn, boundaries, *args):
#   def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):
#   def apply_batch_reshape(self, wrapper, shape):
#   def apply_softmax(self, wrapper):
#   def _set_up_cache(self):
# 
# Path: interval_bound_propagation/src/verifiable_wrapper.py
# class VerifiableWrapper(object):
# class ModelInputWrapper(object):
# class ConstWrapper(VerifiableWrapper):
# class LinearFCWrapper(VerifiableWrapper):
# class LinearConvWrapper(VerifiableWrapper):
# class LinearConv1dWrapper(LinearConvWrapper):
# class LinearConv2dWrapper(LinearConvWrapper):
# class IncreasingMonotonicWrapper(VerifiableWrapper):
# class SoftmaxWrapper(VerifiableWrapper):
# class PiecewiseMonotonicWrapper(VerifiableWrapper):
# class ImageNormWrapper(IncreasingMonotonicWrapper):
# class BatchNormWrapper(VerifiableWrapper):
# class BatchReshapeWrapper(VerifiableWrapper):
# class BatchFlattenWrapper(BatchReshapeWrapper):
#   def __init__(self, module):
#   def input_bounds(self):
#   def output_bounds(self):
#   def module(self):
#   def __str__(self):
#   def propagate_bounds(self, *input_bounds):
#   def _propagate_through(self, module, *input_bounds):
#   def __init__(self, index):
#   def index(self):
#   def output_bounds(self):
#   def output_bounds(self, bounds):
#   def __str__(self):
#   def _propagate_through(self, module):
#   def __init__(self, module):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module, **parameters):
#   def parameters(self):
#   def _propagate_through(self, module, main_bounds, *other_input_bounds):
#   def __init__(self):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module, boundaries=()):
#   def boundaries(self):
#   def _propagate_through(self, module, main_bounds, *other_input_bounds):
#   def __init__(self, module):
#   def inner_module(self):
#   def __init__(self, module):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module, shape):
#   def shape(self):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module):


import abc
import six
import sonnet as snt
import tensorflow.compat.v1 as tf
from absl import logging
from interval_bound_propagation.src import bounds as bounds_lib
from interval_bound_propagation.src import verifiable_wrapper
      num_outputs = tf.shape(c)[2]
      irrelevant = tf.tile(irrelevant, [1, 1, num_outputs])
      self._c = tf.reshape(
          tf.boolean_mask(c, tf.logical_not(irrelevant)),
          [batch_size, -1, num_outputs])
    else:
      self._c = c
    self._d = d

  def _build(self, modules):
    """Outputs specification value."""
    # inputs have shape [batch_size, num_outputs].
    if not (self.collapse and
            isinstance(modules[-1], verifiable_wrapper.LinearFCWrapper)):
      logging.info('Elision of last layer disabled.')
      bounds = modules[-1].output_bounds
      w = self._c
      b = self._d
    else:
      logging.info('Elision of last layer active.')
      # Collapse the last layer.
      bounds = modules[-1].input_bounds
      w = modules[-1].module.w
      b = modules[-1].module.b
      w = tf.einsum('ijk,lk->ijl', self._c, w)
      b = tf.einsum('ijk,k->ij', self._c, b)
      if self._d is not None:
        b += self._d

    # Maximize z * w + b s.t. lower <= z <= upper.

# Path: interval_bound_propagation/src/bounds.py
# class AbstractBounds(object):
# class IntervalBounds(AbstractBounds):
#   def __init__(self):
#   def convert(cls, bounds):
#   def shape(self):
#   def concretize(self):
#   def _raise_not_implemented(self, name):
#   def apply_linear(self, wrapper, w, b):  # pylint: disable=unused-argument
#   def apply_conv1d(self, wrapper, w, b, padding, stride):  # pylint: disable=unused-argument
#   def apply_conv2d(self, wrapper, w, b, padding, strides):  # pylint: disable=unused-argument
#   def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):  # pylint: disable=unused-argument
#   def apply_piecewise_monotonic_fn(self, wrapper, fn, boundaries, *args):  # pylint: disable=unused-argument
#   def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):  # pylint: disable=unused-argument
#   def apply_batch_reshape(self, wrapper, shape):  # pylint: disable=unused-argument
#   def apply_softmax(self, wrapper):  # pylint: disable=unused-argument
#   def update_cache_op(self):
#   def enable_caching(self):
#   def _set_up_cache(self):
#   def _cache_with_update_op(self, tensor):
#   def __init__(self, lower, upper):
#   def lower(self):
#   def upper(self):
#   def shape(self):
#   def __iter__(self):
#   def convert(cls, bounds):
#   def apply_linear(self, wrapper, w, b):
#   def apply_conv1d(self, wrapper, w, b, padding, stride):
#   def apply_conv2d(self, wrapper, w, b, padding, strides):
#   def _affine(self, w, b, fn, **kwargs):
#   def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):
#   def apply_piecewise_monotonic_fn(self, wrapper, fn, boundaries, *args):
#   def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):
#   def apply_batch_reshape(self, wrapper, shape):
#   def apply_softmax(self, wrapper):
#   def _set_up_cache(self):
# 
# Path: interval_bound_propagation/src/verifiable_wrapper.py
# class VerifiableWrapper(object):
# class ModelInputWrapper(object):
# class ConstWrapper(VerifiableWrapper):
# class LinearFCWrapper(VerifiableWrapper):
# class LinearConvWrapper(VerifiableWrapper):
# class LinearConv1dWrapper(LinearConvWrapper):
# class LinearConv2dWrapper(LinearConvWrapper):
# class IncreasingMonotonicWrapper(VerifiableWrapper):
# class SoftmaxWrapper(VerifiableWrapper):
# class PiecewiseMonotonicWrapper(VerifiableWrapper):
# class ImageNormWrapper(IncreasingMonotonicWrapper):
# class BatchNormWrapper(VerifiableWrapper):
# class BatchReshapeWrapper(VerifiableWrapper):
# class BatchFlattenWrapper(BatchReshapeWrapper):
#   def __init__(self, module):
#   def input_bounds(self):
#   def output_bounds(self):
#   def module(self):
#   def __str__(self):
#   def propagate_bounds(self, *input_bounds):
#   def _propagate_through(self, module, *input_bounds):
#   def __init__(self, index):
#   def index(self):
#   def output_bounds(self):
#   def output_bounds(self, bounds):
#   def __str__(self):
#   def _propagate_through(self, module):
#   def __init__(self, module):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module, **parameters):
#   def parameters(self):
#   def _propagate_through(self, module, main_bounds, *other_input_bounds):
#   def __init__(self):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module, boundaries=()):
#   def boundaries(self):
#   def _propagate_through(self, module, main_bounds, *other_input_bounds):
#   def __init__(self, module):
#   def inner_module(self):
#   def __init__(self, module):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module, shape):
#   def shape(self):
#   def _propagate_through(self, module, input_bounds):
#   def __init__(self, module):


# Path: interval_bound_propagation/src/specification.py
import abc
import six
import sonnet as snt
import tensorflow.compat.v1 as tf
from absl import logging
from interval_bound_propagation.src import bounds as bounds_lib
from interval_bound_propagation.src import verifiable_wrapper

      num_outputs = tf.shape(c)[2]
      irrelevant = tf.tile(irrelevant, [1, 1, num_outputs])
      self._c = tf.reshape(
          tf.boolean_mask(c, tf.logical_not(irrelevant)),
          [batch_size, -1, num_outputs])
    else:
      self._c = c
    self._d = d

  def _build(self, modules):
    """Outputs specification value."""
    # inputs have shape [batch_size, num_outputs].
    if not (self.collapse and
            isinstance(modules[-1], verifiable_wrapper.LinearFCWrapper)):
      logging.info('Elision of last layer disabled.')
      bounds = modules[-1].output_bounds
      w = self._c
      b = self._d
    else:
      logging.info('Elision of last layer active.')
      # Collapse the last layer.
      bounds = modules[-1].input_bounds
      w = modules[-1].module.w
      b = modules[-1].module.b
      w = tf.einsum('ijk,lk->ijl', self._c, w)
      b = tf.einsum('ijk,k->ij', self._c, b)
      if self._d is not None:
        b += self._d

    # Maximize z * w + b s.t. lower <= z <= upper.

    bounds = bounds_lib.IntervalBounds.convert(bounds)