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


from interval_bound_propagation.src import layers
import sonnet as snt
import tensorflow.compat.v1 as tf
       output_length,
       input_channels, output_channels]))
  w = tf.transpose(w, perm=[0, 2, 1, 3])
  w = tf.reshape(w, shape=(
      [input_length * input_channels,
       output_length * output_channels]))

  # Broadcast b over spatial dimensions.
  b = tf.tile(b, [output_length]) if b is not None else None

  return w, b


def decode_batchnorm(batchnorm_module):
  """Calculates the neuron-wise multipliers and biases of the batch norm layer.

  Note that, in the case of a convolution, the returned bias will have
  spatial dimensions.

  Args:
    batchnorm_module: `snt.BatchNorm` module.

  Returns:
    w: 1D tensor of shape (output_size) or 3D tensor of shape
      (output_height, output_width, output_channels) containing
      neuron-wise multipliers for the batch norm layer.
    b: 1D tensor of shape (output_size) or 3D tensor of shape
      (output_height, output_width, output_channels) containing
      neuron-wise biases for the batch norm layer.
  """

# 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/layer_utils.py
from interval_bound_propagation.src import layers
import sonnet as snt
import tensorflow.compat.v1 as tf

       output_length,
       input_channels, output_channels]))
  w = tf.transpose(w, perm=[0, 2, 1, 3])
  w = tf.reshape(w, shape=(
      [input_length * input_channels,
       output_length * output_channels]))

  # Broadcast b over spatial dimensions.
  b = tf.tile(b, [output_length]) if b is not None else None

  return w, b


def decode_batchnorm(batchnorm_module):
  """Calculates the neuron-wise multipliers and biases of the batch norm layer.

  Note that, in the case of a convolution, the returned bias will have
  spatial dimensions.

  Args:
    batchnorm_module: `snt.BatchNorm` module.

  Returns:
    w: 1D tensor of shape (output_size) or 3D tensor of shape
      (output_height, output_width, output_channels) containing
      neuron-wise multipliers for the batch norm layer.
    b: 1D tensor of shape (output_size) or 3D tensor of shape
      (output_height, output_width, output_channels) containing
      neuron-wise biases for the batch norm layer.
  """

  if isinstance(batchnorm_module, layers.BatchNorm):