# 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

class LinearSpecification(Specification):
  """Linear specifications: c^T * z_K + d <= 0."""

  def __init__(self, c, d=None, prune_irrelevant=True, collapse=True):
    """Builds a linear specification module."""
    super(LinearSpecification, self).__init__(name='specs', collapse=collapse)
    # c has shape [batch_size, num_specifications, num_outputs]
    # d has shape [batch_size, num_specifications]
    # Some specifications may be irrelevant (not a function of the output).
    # We automatically remove them for clarity. We expect the number of
    # irrelevant specs to be equal for all elements of a batch.
    # Shape is [batch_size, num_specifications]
    if prune_irrelevant:
      irrelevant = tf.equal(tf.reduce_sum(
          tf.cast(tf.abs(c) > 1e-6, tf.int32), axis=-1, keepdims=True), 0)
      batch_size = tf.shape(c)[0]
      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

# 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


class LinearSpecification(Specification):
  """Linear specifications: c^T * z_K + d <= 0."""

  def __init__(self, c, d=None, prune_irrelevant=True, collapse=True):
    """Builds a linear specification module."""
    super(LinearSpecification, self).__init__(name='specs', collapse=collapse)
    # c has shape [batch_size, num_specifications, num_outputs]
    # d has shape [batch_size, num_specifications]
    # Some specifications may be irrelevant (not a function of the output).
    # We automatically remove them for clarity. We expect the number of
    # irrelevant specs to be equal for all elements of a batch.
    # Shape is [batch_size, num_specifications]
    if prune_irrelevant:
      irrelevant = tf.equal(tf.reduce_sum(
          tf.cast(tf.abs(c) > 1e-6, tf.int32), axis=-1, keepdims=True), 0)
      batch_size = tf.shape(c)[0]
      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)):