# Path: interval_bound_propagation/src/layers.py
# class BatchNorm(snt.BatchNorm):
# class ImageNorm(snt.AbstractModule):
#   def __init__(self, axis=None, offset=True, scale=False,
#                decay_rate=0.999, eps=1e-3, initializers=None,
#                partitioners=None, regularizers=None,
#                update_ops_collection=None, name='batch_norm'):
#   def _build_statistics(self, input_batch, axis, use_batch_stats, stat_dtype):
#   def _build(self, input_batch, is_training=True, test_local_stats=False,
#              reuse=False):
#   def scale(self):
#   def bias(self):
#   def mean(self):
#   def variance(self):
#   def epsilon(self):
#   def __init__(self, mean, std, name='image_norm'):
#   def _build(self, inputs):
#   def scale(self):
#   def offset(self):
#   def apply(self, inputs):


import abc
import types
import six
import sonnet as snt
import tensorflow.compat.v1 as tf
from absl import logging
from interval_bound_propagation.src import layers
  """Wraps softmax layers."""

  def __init__(self):
    super(SoftmaxWrapper, self).__init__(None)

  def _propagate_through(self, module, input_bounds):
    return input_bounds.apply_softmax(self)


class PiecewiseMonotonicWrapper(VerifiableWrapper):
  """Wraps a piecewise (not necessarily increasing) monotonic function."""

  def __init__(self, module, boundaries=()):
    super(PiecewiseMonotonicWrapper, self).__init__(module)
    self._boundaries = boundaries

  @property
  def boundaries(self):
    return self._boundaries

  def _propagate_through(self, module, main_bounds, *other_input_bounds):
    return main_bounds.apply_piecewise_monotonic_fn(self, module,
                                                    self.boundaries,
                                                    *other_input_bounds)


class ImageNormWrapper(IncreasingMonotonicWrapper):
  """Convenience wrapper for getting track of the ImageNorm layer."""

  def __init__(self, module):

# Path: interval_bound_propagation/src/layers.py
# class BatchNorm(snt.BatchNorm):
# class ImageNorm(snt.AbstractModule):
#   def __init__(self, axis=None, offset=True, scale=False,
#                decay_rate=0.999, eps=1e-3, initializers=None,
#                partitioners=None, regularizers=None,
#                update_ops_collection=None, name='batch_norm'):
#   def _build_statistics(self, input_batch, axis, use_batch_stats, stat_dtype):
#   def _build(self, input_batch, is_training=True, test_local_stats=False,
#              reuse=False):
#   def scale(self):
#   def bias(self):
#   def mean(self):
#   def variance(self):
#   def epsilon(self):
#   def __init__(self, mean, std, name='image_norm'):
#   def _build(self, inputs):
#   def scale(self):
#   def offset(self):
#   def apply(self, inputs):


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

  """Wraps softmax layers."""

  def __init__(self):
    super(SoftmaxWrapper, self).__init__(None)

  def _propagate_through(self, module, input_bounds):
    return input_bounds.apply_softmax(self)


class PiecewiseMonotonicWrapper(VerifiableWrapper):
  """Wraps a piecewise (not necessarily increasing) monotonic function."""

  def __init__(self, module, boundaries=()):
    super(PiecewiseMonotonicWrapper, self).__init__(module)
    self._boundaries = boundaries

  @property
  def boundaries(self):
    return self._boundaries

  def _propagate_through(self, module, main_bounds, *other_input_bounds):
    return main_bounds.apply_piecewise_monotonic_fn(self, module,
                                                    self.boundaries,
                                                    *other_input_bounds)


class ImageNormWrapper(IncreasingMonotonicWrapper):
  """Convenience wrapper for getting track of the ImageNorm layer."""

  def __init__(self, module):

    if not isinstance(module, layers.ImageNorm):