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


from interval_bound_propagation.src import bounds as basic_bounds
from interval_bound_propagation.src import relative_bounds
import sonnet as snt
import tensorflow.compat.v1 as tf
  def shape(self):
    return self.nominal.shape.as_list()

  @classmethod
  def convert(cls, bounds):
    if not isinstance(bounds, cls):
      raise ValueError('Cannot convert "{}" to "{}"'.format(bounds,
                                                            cls.__name__))
    return bounds

  def apply_batch_reshape(self, wrapper, shape):
    reshape = snt.BatchReshape(shape)
    if self.vertices.shape.ndims == self.nominal.shape.ndims:
      reshape_vertices = reshape
    else:
      reshape_vertices = snt.BatchReshape(shape, preserve_dims=2)
    return SimplexBounds(reshape_vertices(self.vertices),
                         reshape(self.nominal),
                         self.r)

  def apply_linear(self, wrapper, w, b):
    mapped_centres = tf.matmul(self.nominal, w)
    mapped_vertices = tf.tensordot(self.vertices, w, axes=1)

    lb, ub = _simplex_bounds(mapped_vertices, mapped_centres, self.r, -2)

    nominal_out = tf.matmul(self.nominal, w)
    if b is not None:
      nominal_out += b


# 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/simplex_bounds.py
from interval_bound_propagation.src import bounds as basic_bounds
from interval_bound_propagation.src import relative_bounds
import sonnet as snt
import tensorflow.compat.v1 as tf

  def shape(self):
    return self.nominal.shape.as_list()

  @classmethod
  def convert(cls, bounds):
    if not isinstance(bounds, cls):
      raise ValueError('Cannot convert "{}" to "{}"'.format(bounds,
                                                            cls.__name__))
    return bounds

  def apply_batch_reshape(self, wrapper, shape):
    reshape = snt.BatchReshape(shape)
    if self.vertices.shape.ndims == self.nominal.shape.ndims:
      reshape_vertices = reshape
    else:
      reshape_vertices = snt.BatchReshape(shape, preserve_dims=2)
    return SimplexBounds(reshape_vertices(self.vertices),
                         reshape(self.nominal),
                         self.r)

  def apply_linear(self, wrapper, w, b):
    mapped_centres = tf.matmul(self.nominal, w)
    mapped_vertices = tf.tensordot(self.vertices, w, axes=1)

    lb, ub = _simplex_bounds(mapped_vertices, mapped_centres, self.r, -2)

    nominal_out = tf.matmul(self.nominal, w)
    if b is not None:
      nominal_out += b


    return relative_bounds.RelativeIntervalBounds(lb, ub, nominal_out)