# 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/relative_bounds.py
# class RelativeIntervalBounds(basic_bounds.AbstractBounds):
#   def __init__(self, lower_offset, upper_offset, nominal):
#   def lower_offset(self):
#   def upper_offset(self):
#   def nominal(self):
#   def lower(self):
#   def upper(self):
#   def shape(self):
#   def convert(cls, bounds):
#   def apply_batch_reshape(self, wrapper, shape):
#   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 apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):
#   def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):
#   def _set_up_cache(self):
# def _maxpool_bounds(module, kernel_shape, strides, lb_in, ub_in,
#                     nominal_in, nominal_out):
# def _activation_bounds(nl_fun, lb_in, ub_in, nominal_in, parameters=None):


import collections
import sonnet as snt
import tensorflow.compat.v1 as tf
from absl import logging
from interval_bound_propagation.src import bounds as basic_bounds
from interval_bound_propagation.src import relative_bounds
    zero = tf.zeros_like(lb)
    lb = snt.BatchFlatten()(lb)
    ub = snt.BatchFlatten()(ub)
    input_size = tf.shape(lb)[1]
    output_shape = tf.concat([[input_size], input_shape], axis=0)
    identity = tf.reshape(tf.eye(input_size), output_shape)
    identity = tf.expand_dims(identity, 0)
    identity = tf.tile(identity, [batch_size] + [1] * (len(input_shape) + 1))
    expr = LinearExpression(w=identity, b=zero,
                            lower=lb, upper=ub)
    return expr, expr


class RelativeSymbolicBounds(SymbolicBounds):
  """Relative-to-nominal variant of Fast-Lin bounds."""

  def __init__(self, lower_offset, upper_offset, nominal):
    super(RelativeSymbolicBounds, self).__init__(lower_offset, upper_offset)
    self._nominal = nominal

  def concretize(self):
    """Returns lower and upper interval bounds."""
    if self._concretized is None:
      # Construct once and cache.
      lb_offset, ub_offset = self._concretize_bounds(self.lower, self.upper)

      # Apply intersections with prior runs.
      if self._prior_bounds is not None:
        lb_offset = tf.maximum(lb_offset, self._prior_bounds.lower_offset)
        ub_offset = tf.minimum(ub_offset, self._prior_bounds.upper_offset)

# 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/relative_bounds.py
# class RelativeIntervalBounds(basic_bounds.AbstractBounds):
#   def __init__(self, lower_offset, upper_offset, nominal):
#   def lower_offset(self):
#   def upper_offset(self):
#   def nominal(self):
#   def lower(self):
#   def upper(self):
#   def shape(self):
#   def convert(cls, bounds):
#   def apply_batch_reshape(self, wrapper, shape):
#   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 apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):
#   def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):
#   def _set_up_cache(self):
# def _maxpool_bounds(module, kernel_shape, strides, lb_in, ub_in,
#                     nominal_in, nominal_out):
# def _activation_bounds(nl_fun, lb_in, ub_in, nominal_in, parameters=None):


# Path: interval_bound_propagation/src/fastlin.py
import collections
import sonnet as snt
import tensorflow.compat.v1 as tf
from absl import logging
from interval_bound_propagation.src import bounds as basic_bounds
from interval_bound_propagation.src import relative_bounds

    zero = tf.zeros_like(lb)
    lb = snt.BatchFlatten()(lb)
    ub = snt.BatchFlatten()(ub)
    input_size = tf.shape(lb)[1]
    output_shape = tf.concat([[input_size], input_shape], axis=0)
    identity = tf.reshape(tf.eye(input_size), output_shape)
    identity = tf.expand_dims(identity, 0)
    identity = tf.tile(identity, [batch_size] + [1] * (len(input_shape) + 1))
    expr = LinearExpression(w=identity, b=zero,
                            lower=lb, upper=ub)
    return expr, expr


class RelativeSymbolicBounds(SymbolicBounds):
  """Relative-to-nominal variant of Fast-Lin bounds."""

  def __init__(self, lower_offset, upper_offset, nominal):
    super(RelativeSymbolicBounds, self).__init__(lower_offset, upper_offset)
    self._nominal = nominal

  def concretize(self):
    """Returns lower and upper interval bounds."""
    if self._concretized is None:
      # Construct once and cache.
      lb_offset, ub_offset = self._concretize_bounds(self.lower, self.upper)

      # Apply intersections with prior runs.
      if self._prior_bounds is not None:
        lb_offset = tf.maximum(lb_offset, self._prior_bounds.lower_offset)
        ub_offset = tf.minimum(ub_offset, self._prior_bounds.upper_offset)

      self._concretized = relative_bounds.RelativeIntervalBounds(