class Clip(Constraint):
  """Clips weight constraint."""

  # This function was modified from Keras minmaxconstraints.
  #
  # Constrains the weights to be between min/max values.
  #   min_value: the minimum norm for the incoming weights.
  #   max_value: the maximum norm for the incoming weights.
  #   constraint: previous constraint to be clipped.
  #   quantizer: quantizer to be applied to constraint.

  def __init__(self, min_value=0.0, max_value=1.0,
               constraint=None, quantizer=None):
    """Initializes Clip constraint class."""

    self.min_value = min_value
    self.max_value = max_value
    self.constraint = constraints.get(constraint)
    # Don't wrap yourself
    if isinstance(self.constraint, Clip):
      self.constraint = None
    self.quantizer = get_quantizer(quantizer)

  def __call__(self, w):
    """Clips values between min and max values."""
    if self.constraint:
      w = self.constraint(w)
      if self.quantizer:
        w = self.quantizer(w)
    w = tf.keras.backend.clip(w, self.min_value, self.max_value)
    return w

  def get_config(self):
    """Returns configuration of constraint class."""
    return {"min_value": self.min_value, "max_value": self.max_value}

  @classmethod
  def from_config(cls, config):
    if isinstance(config.get('constraint', None), Clip):
      config['constraint'] = None
    config['constraint'] = constraints.get(config.get('constraint', None))
    config['quantizer'] = get_quantizer(config.get('quantizer', None))
    return cls(**config) class QConv2DBatchnorm(QConv2D):
  """Fold batchnormalization with a previous qconv2d layer."""

  def __init__(
      self,
      # qconv2d params
      filters,
      kernel_size,
      strides=(1, 1),
      padding="valid",
      data_format="channels_last",
      dilation_rate=(1, 1),
      activation=None,
      use_bias=True,
      kernel_initializer="he_normal",
      bias_initializer="zeros",
      kernel_regularizer=None,
      bias_regularizer=None,
      activity_regularizer=None,
      kernel_constraint=None,
      bias_constraint=None,
      kernel_quantizer=None,
      bias_quantizer=None,

      # batchnorm params
      axis=-1,
      momentum=0.99,
      epsilon=0.001,
      center=True,
      scale=True,
      beta_initializer="zeros",
      gamma_initializer="ones",
      moving_mean_initializer="zeros",
      moving_variance_initializer="ones",
      beta_regularizer=None,
      gamma_regularizer=None,
      beta_constraint=None,
      gamma_constraint=None,
      renorm=False,
      renorm_clipping=None,
      renorm_momentum=0.99,
      fused=None,
      trainable=True,
      virtual_batch_size=None,
      adjustment=None,

      # other params
      ema_freeze_delay=None,
      folding_mode="ema_stats_folding",
      **kwargs):
    """Initialize a composite layer that folds conv2d and batch normalization.

    The first group of parameters correponds to the initialization parameters
      of a qconv2d layer. check qkeras.qconvolutional.qconv2d for details.

    The 2nd group of parameters corresponds to the initialization parameters
      of a BatchNormalization layer. Check keras.layers.normalization.BatchNorma
      lizationBase for details.

    The 3rd group of parameters corresponds to the initialization parameters
      specific to this class.

      ema_freeze_delay: int. number of steps before batch normalization mv_mean
        and mv_variance will be frozen and used in the folded layer.
      folding_mode: string
        "ema_stats_folding": mimic tflite which uses the ema statistics to
          fold the kernel to suppress quantization induced jitter then performs
          the correction to have a similar effect of using the current batch
          statistics.
        "batch_stats_folding": use batch mean and variance to fold kernel first;
          after enough training steps switch to moving_mean and moving_variance
          for kernel folding.
    """

    # intialization the qconv2d part of the composite layer
    super(QConv2DBatchnorm, self).__init__(
        filters=filters,
        kernel_size=kernel_size,
        strides=strides,
        padding=padding,
        dilation_rate=dilation_rate,
        activation=activation,
        use_bias=use_bias,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        kernel_constraint=kernel_constraint,
        bias_constraint=bias_constraint,
        kernel_quantizer=kernel_quantizer,
        bias_quantizer=bias_quantizer,
        **kwargs)

    # initialization of batchnorm part of the composite layer
    self.batchnorm = layers.BatchNormalization(
        axis=axis, momentum=momentum, epsilon=epsilon, center=center,
        scale=scale, beta_initializer=beta_initializer,
        gamma_initializer=gamma_initializer,
        moving_mean_initializer=moving_mean_initializer,
        moving_variance_initializer=moving_variance_initializer,
        beta_regularizer=beta_regularizer,
        gamma_regularizer=gamma_regularizer,
        beta_constraint=beta_constraint, gamma_constraint=gamma_constraint,
        renorm=renorm, renorm_clipping=renorm_clipping, 
        renorm_momentum=renorm_momentum, fused=fused, trainable=trainable,
        virtual_batch_size=virtual_batch_size, adjustment=adjustment)

    self.ema_freeze_delay = ema_freeze_delay
    assert folding_mode in ["ema_stats_folding", "batch_stats_folding"]
    self.folding_mode = folding_mode

  def build(self, input_shape):
    super(QConv2DBatchnorm, self).build(input_shape)

    # self._iteration (i.e., training_steps) is initialized with -1. When
    # loading ckpt, it can load the number of training steps that have been
    # previously trainied. If start training from scratch.
    # TODO(lishanok): develop a way to count iterations outside layer
    self._iteration = tf.Variable(-1, trainable=False, name="iteration",
                                  dtype=tf.int64)

  def call(self, inputs, training=None):

    # numpy value, mark the layer is in training
    training = self.batchnorm._get_training_value(training)  # pylint: disable=protected-access

    # checking if to update batchnorm params
    if (self.ema_freeze_delay is None) or (self.ema_freeze_delay < 0):
      # if ema_freeze_delay is None or a negative value, do not freeze bn stats
      bn_training = tf.cast(training, dtype=bool)
    else:
      bn_training = tf.math.logical_and(training, tf.math.less_equal(
          self._iteration, self.ema_freeze_delay))

    kernel = self.kernel

    # run conv to produce output for the following batchnorm
    conv_outputs = tf.keras.backend.conv2d(
        inputs,
        kernel,
        strides=self.strides,
        padding=self.padding,
        data_format=self.data_format,
        dilation_rate=self.dilation_rate)

    if self.use_bias:
      bias = self.bias
      conv_outputs = tf.keras.backend.bias_add(
          conv_outputs, bias, data_format=self.data_format)
    else:
      bias = 0

    _ = self.batchnorm(conv_outputs, training=bn_training)

    self._iteration.assign_add(tf_utils.smart_cond(
        training, lambda: tf.constant(1, tf.int64),
        lambda: tf.constant(0, tf.int64)))

    # calcuate mean and variance from current batch
    bn_shape = conv_outputs.shape
    ndims = len(bn_shape)
    reduction_axes = [i for i in range(ndims) if i not in self.batchnorm.axis]
    keep_dims = len(self.batchnorm.axis) > 1
    mean, variance = self.batchnorm._moments(  # pylint: disable=protected-access
        math_ops.cast(conv_outputs, self.batchnorm._param_dtype),  # pylint: disable=protected-access
        reduction_axes,
        keep_dims=keep_dims)
    # get batchnorm weights
    gamma = self.batchnorm.gamma
    beta = self.batchnorm.beta
    moving_mean = self.batchnorm.moving_mean
    moving_variance = self.batchnorm.moving_variance

    if self.folding_mode == "batch_stats_folding":
      # using batch mean and variance in the initial training stage
      # after sufficient training, switch to moving mean and variance
      new_mean = tf_utils.smart_cond(
          bn_training, lambda: mean, lambda: moving_mean)
      new_variance = tf_utils.smart_cond(
          bn_training, lambda: variance, lambda: moving_variance)

      # get the inversion factor so that we replace division by multiplication
      inv = math_ops.rsqrt(new_variance + self.batchnorm.epsilon)
      if gamma is not None:
        inv *= gamma
      # fold bias with bn stats
      folded_bias = inv * (bias - new_mean) + beta

    elif self.folding_mode == "ema_stats_folding":
      # We always scale the weights with a correction factor to the long term
      # statistics prior to quantization. This ensures that there is no jitter
      # in the quantized weights due to batch to batch variation. During the
      # initial phase of training, we undo the scaling of the weights so that
      # outputs are identical to regular batch normalization. We also modify
      # the bias terms correspondingly. After sufficient training, switch from
      # using batch statistics to long term moving averages for batch
      # normalization.

      # use batch stats for calcuating bias before bn freeze, and use moving
      # stats after bn freeze
      mv_inv = math_ops.rsqrt(moving_variance + self.batchnorm.epsilon)
      batch_inv = math_ops.rsqrt(variance + self.batchnorm.epsilon)

      if gamma is not None:
        mv_inv *= gamma
        batch_inv *= gamma
      folded_bias = tf_utils.smart_cond(
          bn_training,
          lambda: batch_inv * (bias - mean) + beta,
          lambda: mv_inv * (bias - moving_mean) + beta)
      # moving stats is always used to fold kernel in tflite; before bn freeze
      # an additional correction factor will be applied to the conv2d output
      inv = mv_inv
    else:
      assert ValueError

    # wrap conv kernel with bn parameters
    folded_kernel = inv * kernel
    # quantize the folded kernel
    if self.kernel_quantizer is not None:
      q_folded_kernel = self.kernel_quantizer_internal(folded_kernel)
    else:
      q_folded_kernel = folded_kernel

    # If loaded from a ckpt, bias_quantizer is the ckpt value
    # Else if bias_quantizer not specified, bias
    #   quantizer is None and we need to calculate bias quantizer
    #   type according to accumulator type. User can call
    #   bn_folding_utils.populate_bias_quantizer_from_accumulator(
    #      model, input_quantizer_list]) to populate such bias quantizer.
    if self.bias_quantizer_internal is not None:
      q_folded_bias = self.bias_quantizer_internal(folded_bias)
    else:
      q_folded_bias = folded_bias

    applied_kernel = q_folded_kernel
    applied_bias = q_folded_bias

    # calculate conv2d output using the quantized folded kernel
    folded_outputs = tf.keras.backend.conv2d(
        inputs,
        applied_kernel,
        strides=self.strides,
        padding=self.padding,
        data_format=self.data_format,
        dilation_rate=self.dilation_rate)
    if training is True and self.folding_mode == "ema_stats_folding":
      batch_inv = math_ops.rsqrt(variance + self.batchnorm.epsilon)
      y_corr = tf_utils.smart_cond(
          bn_training,
          lambda: (math_ops.sqrt(moving_variance + self.batchnorm.epsilon) *
                   math_ops.rsqrt(variance + self.batchnorm.epsilon)),
          lambda: tf.constant(1.0, shape=moving_variance.shape))
      folded_outputs = math_ops.mul(folded_outputs, y_corr)

    folded_outputs = tf.keras.backend.bias_add(
        folded_outputs,
        applied_bias,
        data_format=self.data_format)
    if self.activation is not None:
      return self.activation(folded_outputs)

    return folded_outputs

  def get_config(self):
    base_config = super().get_config()
    bn_config = self.batchnorm.get_config()
    config = {"ema_freeze_delay": self.ema_freeze_delay,
              "folding_mode": self.folding_mode}
    name = base_config["name"]
    out_config = dict(
        list(base_config.items())
        + list(bn_config.items()) + list(config.items()))

    # names from different config override each other; use the base layer name
    # as the this layer's config name
    out_config["name"] = name
    return out_config

  def get_quantization_config(self):
    return {
        "kernel_quantizer": str(self.kernel_quantizer_internal),
        "bias_quantizer": str(self.bias_quantizer_internal),
        "activation": str(self.activation),
        "filters": str(self.filters)
    }

  def get_quantizers(self):
    return self.quantizers

  def get_folded_weights(self):
    """Function to get the batchnorm folded weights.

    This function converts the weights by folding batchnorm parameters into
    the weight of QConv2D. The high-level equation:

    W_fold = gamma * W / sqrt(variance + epsilon)
    bias_fold = gamma * (bias - moving_mean) / sqrt(variance + epsilon) + beta
    """

    kernel = self.kernel
    if self.use_bias:
      bias = self.bias
    else:
      bias = 0

    # get batchnorm weights and moving stats
    gamma = self.batchnorm.gamma
    beta = self.batchnorm.beta
    moving_mean = self.batchnorm.moving_mean
    moving_variance = self.batchnorm.moving_variance
    # get the inversion factor so that we replace division by multiplication
    inv = math_ops.rsqrt(moving_variance + self.batchnorm.epsilon)
    if gamma is not None:
      inv *= gamma

    # wrap conv kernel and bias with bn parameters
    folded_kernel = inv * kernel
    folded_bias = inv * (bias - moving_mean) + beta

    return [folded_kernel, folded_bias] class QDepthwiseConv2DBatchnorm(QDepthwiseConv2D):
  """Fold batchnormalization with a previous QDepthwiseConv2d layer."""

  def __init__(
      self,
      # QDepthwiseConv2d params
      kernel_size,
      strides=(1, 1),
      padding="VALID",
      depth_multiplier=1,
      data_format=None,
      activation=None,
      use_bias=True,
      depthwise_initializer="he_normal",
      bias_initializer="zeros",
      depthwise_regularizer=None,
      bias_regularizer=None,
      activity_regularizer=None,
      depthwise_constraint=None,
      bias_constraint=None,
      dilation_rate=(1, 1),
      depthwise_quantizer=None,
      bias_quantizer=None,
      depthwise_range=None,
      bias_range=None,

      # batchnorm params
      axis=-1,
      momentum=0.99,
      epsilon=0.001,
      center=True,
      scale=True,
      beta_initializer="zeros",
      gamma_initializer="ones",
      moving_mean_initializer="zeros",
      moving_variance_initializer="ones",
      beta_regularizer=None,
      gamma_regularizer=None,
      beta_constraint=None,
      gamma_constraint=None,
      renorm=False,
      renorm_clipping=None,
      renorm_momentum=0.99,
      fused=None,
      trainable=True,
      virtual_batch_size=None,
      adjustment=None,

      # other params
      ema_freeze_delay=None,
      folding_mode="ema_stats_folding",
      **kwargs):

    """A composite layer that folds depthwiseconv2d and batch normalization.

    The first group of parameters correponds to the initialization parameters
      of a QDepthwiseConv2d layer. check qkeras.qconvolutional.QDepthwiseConv2D
      for details.

    The 2nd group of parameters corresponds to the initialization parameters
      of a BatchNormalization layer. Check keras.layers.normalization.BatchNorma
      lizationBase for details.

    The 3rd group of parameters corresponds to the initialization parameters
      specific to this class.

      ema_freeze_delay: int or None. number of steps before batch normalization
        mv_mean and mv_variance will be frozen and used in the folded layer.
      folding_mode: string
        "ema_stats_folding": mimic tflite which uses the ema statistics to
          fold the kernel to suppress quantization induced jitter then performs
          the correction to have a similar effect of using the current batch
          statistics.
        "batch_stats_folding": use batch mean and variance to fold kernel first;
          after enough training steps switch to moving_mean and moving_variance
          for kernel folding.
    """

    # intialization the QDepthwiseConv2d part of the composite layer
    super(QDepthwiseConv2DBatchnorm, self).__init__(
        kernel_size=kernel_size,
        strides=strides,
        padding=padding,
        depth_multiplier=depth_multiplier,
        data_format=data_format,
        activation=activation,
        use_bias=use_bias,
        depthwise_initializer=depthwise_initializer,
        bias_initializer=bias_initializer,
        depthwise_regularizer=depthwise_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        depthwise_constraint=depthwise_constraint,
        bias_constraint=bias_constraint,
        dilation_rate=dilation_rate,
        depthwise_quantizer=depthwise_quantizer,
        bias_quantizer=bias_quantizer,
        depthwise_range=depthwise_range,
        bias_range=bias_range,
        **kwargs)

    # initialization of batchnorm part of the composite layer
    self.batchnorm = layers.BatchNormalization(
        axis=axis, momentum=momentum, epsilon=epsilon, center=center,
        scale=scale, beta_initializer=beta_initializer,
        gamma_initializer=gamma_initializer,
        moving_mean_initializer=moving_mean_initializer,
        moving_variance_initializer=moving_variance_initializer,
        beta_regularizer=beta_regularizer,
        gamma_regularizer=gamma_regularizer,
        beta_constraint=beta_constraint, gamma_constraint=gamma_constraint,
        renorm=renorm, renorm_clipping=renorm_clipping, 
        renorm_momentum=renorm_momentum, fused=fused, trainable=trainable,
        virtual_batch_size=virtual_batch_size, adjustment=adjustment)

    self.ema_freeze_delay = ema_freeze_delay
    assert folding_mode in ["ema_stats_folding", "batch_stats_folding"]
    self.folding_mode = folding_mode

  def build(self, input_shape):
    super(QDepthwiseConv2DBatchnorm, self).build(input_shape)

    # If start training from scratch, self._iteration (i.e., training_steps)
    # is initialized with -1. When loading ckpt, it can load the number of
    # training steps that have been previously trainied.
    # TODO(lishanok): develop a way to count iterations outside layer
    self._iteration = tf.Variable(-1, trainable=False, name="iteration",
                                  dtype=tf.int64)

  def call(self, inputs, training=None):

    # numpy value, mark the layer is in training
    training = self.batchnorm._get_training_value(training)  # pylint: disable=protected-access

    # checking if to update batchnorm params
    if (self.ema_freeze_delay is None) or (self.ema_freeze_delay < 0):
      # if ema_freeze_delay is None or a negative value, do not freeze bn stats
      bn_training = tf.cast(training, dtype=bool)
    else:
      bn_training = tf.math.logical_and(training, tf.math.less_equal(
          self._iteration, self.ema_freeze_delay))

    depthwise_kernel = self.depthwise_kernel

    # run depthwise_conv2d to produce output for the following batchnorm
    conv_outputs = tf.keras.backend.depthwise_conv2d(
        inputs,
        depthwise_kernel,
        strides=self.strides,
        padding=self.padding,
        dilation_rate=self.dilation_rate,
        data_format=self.data_format)

    if self.use_bias:
      bias = self.bias
      conv_outputs = tf.keras.backend.bias_add(
          conv_outputs, bias, data_format=self.data_format)
    else:
      bias = 0

    _ = self.batchnorm(conv_outputs, training=bn_training)

    self._iteration.assign_add(tf_utils.smart_cond(
        training, lambda: tf.constant(1, tf.int64),
        lambda: tf.constant(0, tf.int64)))

    # calcuate mean and variance from current batch
    bn_shape = conv_outputs.shape
    ndims = len(bn_shape)
    reduction_axes = [i for i in range(ndims) if i not in self.batchnorm.axis]
    keep_dims = len(self.batchnorm.axis) > 1
    mean, variance = self.batchnorm._moments(  # pylint: disable=protected-access
        math_ops.cast(conv_outputs, self.batchnorm._param_dtype),  # pylint: disable=protected-access
        reduction_axes,
        keep_dims=keep_dims)
    gamma = self.batchnorm.gamma
    beta = self.batchnorm.beta
    moving_mean = self.batchnorm.moving_mean
    moving_variance = self.batchnorm.moving_variance

    if self.folding_mode not in ["batch_stats_folding", "ema_stats_folding"]:
      assert ValueError("mode {} not supported!".format(self.folding_mode))

    mv_inv = math_ops.rsqrt(moving_variance + self.batchnorm.epsilon)
    batch_inv = math_ops.rsqrt(variance + self.batchnorm.epsilon)

    if gamma is not None:
      mv_inv *= gamma
      batch_inv *= gamma

    folded_bias = tf_utils.smart_cond(
        bn_training,
        lambda: batch_inv * (bias - mean) + beta,
        lambda: mv_inv * (bias - moving_mean) + beta)

    if self.folding_mode == "batch_stats_folding":
      # using batch mean and variance in the initial training stage
      # after sufficient training, switch to moving mean and variance
      inv = tf_utils.smart_cond(bn_training, lambda: batch_inv, lambda: mv_inv)

    elif self.folding_mode == "ema_stats_folding":
      # We always scale the weights with a correction factor to the long term
      # statistics prior to quantization. This ensures that there is no jitter
      # in the quantized weights due to batch to batch variation. During the
      # initial phase of training, we undo the scaling of the weights so that
      # outputs are identical to regular batch normalization. We also modify
      # the bias terms correspondingly. After sufficient training, switch from
      # using batch statistics to long term moving averages for batch
      # normalization.

      # use batch stats for calcuating bias before bn freeze, and use moving
      # stats after bn freeze

      # moving stats is always used to fold kernel in tflite; before bn freeze
      # an additional correction factor will be applied to the depthwiseconv2d
      # output
      inv = mv_inv

    # for DepthwiseConv2D inv needs to be broadcasted to the last 2 dimensions
    # of the kernels
    depthwise_weights_shape = [
        depthwise_kernel.get_shape().as_list()[2],
        depthwise_kernel.get_shape().as_list()[3]
    ]
    inv = array_ops.reshape(inv, depthwise_weights_shape)
    # wrap conv kernel with bn parameters
    folded_depthwise_kernel = inv * depthwise_kernel
    # quantize the folded kernel
    if self.depthwise_quantizer is not None:
      q_folded_depthwise_kernel = self.depthwise_quantizer_internal(
          folded_depthwise_kernel)
    else:
      q_folded_depthwise_kernel = folded_depthwise_kernel

    # If loaded from a ckpt, bias_quantizer is the ckpt value
    # Else if bias_quantizer not specified, bias
    #   quantizer is None and we need to calculate bias quantizer
    #   type according to accumulator type. User can call
    #   bn_folding_utils.populate_bias_quantizer_for_folded_layers(
    #      model, input_quantizer_list]) to populate such bias quantizer.
    if self.bias_quantizer is not None:
      q_folded_bias = self.bias_quantizer_internal(folded_bias)
    else:
      q_folded_bias = folded_bias

    applied_kernel = q_folded_depthwise_kernel
    applied_bias = q_folded_bias

    # calculate depthwise_conv2d output using the quantized folded kernel
    folded_outputs = tf.keras.backend.depthwise_conv2d(
        inputs,
        applied_kernel,
        strides=self.strides,
        padding=self.padding,
        dilation_rate=self.dilation_rate,
        data_format=self.data_format)

    if training is True and self.folding_mode == "ema_stats_folding":
      batch_inv = math_ops.rsqrt(variance + self.batchnorm.epsilon)
      y_corr = tf_utils.smart_cond(
          bn_training,
          lambda: (math_ops.sqrt(moving_variance + self.batchnorm.epsilon) *
                   math_ops.rsqrt(variance + self.batchnorm.epsilon)),
          lambda: tf.constant(1.0, shape=moving_variance.shape))
      folded_outputs = math_ops.mul(folded_outputs, y_corr)

    folded_outputs = tf.keras.backend.bias_add(
        folded_outputs,
        applied_bias,
        data_format=self.data_format)

    if self.activation is not None:
      return self.activation(folded_outputs)

    return folded_outputs

  def get_config(self):
    base_config = super().get_config()
    bn_config = self.batchnorm.get_config()
    config = {"ema_freeze_delay": self.ema_freeze_delay,
              "folding_mode": self.folding_mode}
    name = base_config["name"]
    out_config = dict(
        list(base_config.items())
        + list(bn_config.items()) + list(config.items()))

    # names from different config override each other; use the base layer name
    # as the this layer's config name
    out_config["name"] = name
    return out_config

  def get_quantization_config(self):
    return {
        "depthwise_quantizer": str(self.depthwise_quantizer_internal),
        "bias_quantizer": str(self.bias_quantizer_internal),
        "activation": str(self.activation),
        "filters": str(self.filters)
    }

  def get_quantizers(self):
    return self.quantizers

  def get_folded_weights(self):
    """Function to get the batchnorm folded weights.

    This function converts the weights by folding batchnorm parameters into
    the weight of QDepthwiseConv2d. The high-level equation:

    W_fold = gamma * W / sqrt(variance + epsilon)
    bias_fold = gamma * (bias - moving_mean) / sqrt(variance + epsilon) + beta
    """

    depthwise_kernel = self.depthwise_kernel

    if self.use_bias:
      bias = self.bias
    else:
      bias = 0

    # get Batchnorm stats
    gamma = self.batchnorm.gamma
    beta = self.batchnorm.beta
    moving_mean = self.batchnorm.moving_mean
    moving_variance = self.batchnorm.moving_variance

    # get the inversion factor so that we replace division by multiplication
    inv = math_ops.rsqrt(moving_variance + self.batchnorm.epsilon)
    if gamma is not None:
      inv *= gamma
    # fold bias with bn stats
    folded_bias = inv * (bias - moving_mean) + beta

    # for DepthwiseConv2D inv needs to be broadcasted to the last 2 dimensions
    # of the kernels
    depthwise_weights_shape = [
        depthwise_kernel.get_shape().as_list()[2],
        depthwise_kernel.get_shape().as_list()[3]
    ]
    inv = array_ops.reshape(inv, depthwise_weights_shape)
    # wrap conv kernel with bn parameters
    folded_depthwise_kernel = inv * depthwise_kernel

    return [folded_depthwise_kernel, folded_bias] class QActivation(Layer, PrunableLayer):
  """Implements quantized activation layers."""

  def __init__(self, activation, **kwargs):

    super(QActivation, self).__init__(**kwargs)

    self.activation = activation

    if not isinstance(activation, six.string_types):
      self.quantizer = activation
      if hasattr(self.quantizer, "__name__"):
        self.__name__ = self.quantizer.__name__
      elif hasattr(self.quantizer, "name"):
        self.__name__ = self.quantizer.name
      elif hasattr(self.quantizer, "__class__"):
        self.__name__ = self.quantizer.__class__.__name__
      return

    self.__name__ = activation

    try:
      self.quantizer = get_quantizer(activation)
    except KeyError:
      raise ValueError("invalid activation '{}'".format(activation))

  def call(self, inputs):
    return self.quantizer(inputs)

  def get_config(self):
    config = {"activation": self.activation}
    base_config = super(QActivation, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return str(self.activation)

  def compute_output_shape(self, input_shape):
    return input_shape

  def get_prunable_weights(self):
    return [] class QAdaptiveActivation(Layer, PrunableLayer):
  """[EXPERIMENTAL] Implements an adaptive quantized activation layer using EMA.

  This layer calculates an exponential moving average of min and max of the
  activation values to automatically determine the scale (integer bits) of
  the quantizer used in this layer.
  """

  def __init__(self,
               activation,
               total_bits,
               current_step=None,
               symmetric=True,
               quantization_delay=0,
               ema_freeze_delay=None,
               ema_decay=0.9999,
               per_channel=False,
               po2_rounding=False,
               relu_neg_slope=0.0,
               relu_upper_bound=None,
               **kwargs):
    """Initializes this QAdaptiveActivation layer.

    Args:
      activation: Str. The activation quantizer type to use for this activation
        layer, such as 'quantized_relu'. Should be a string with no params.
      total_bits: Int. The total bits that can be used by the quantizer
      current_step: tf.Variable specifying the current step in training.
        You can find this by passing model.optimizer.iterations
        (see tf.keras.optimizers.Optimizer.iterations). If set to None, the
        layer will attempt to estimate the current step itself, but please note
        that this number may not always match the optimizer step.
      symmetric: Bool. If to enforce symmetry about the origin in the quantized
        bit representation of the value. When using linear activation, this
        should be True for best results.
      quantization_delay: Int. How many training steps to wait until quantizing
        the activation values.
      ema_freeze_delay: Int. Steps to wait until stopping the update of the
        exponential moving average values. Set to None for an infinite delay.
      ema_decay: Float. The decay value used for exponential moving average (see
        tf.keras.backend.moving_average_update)
      per_channel: Bool. If to quantize the activation values on a
        per-channel basis.
      po2_rounding: Bool. If true, the EMA max value is rounded to the nearest
        power-of-2. If false, the EMA max value is rounded up (with ceil) to a
        power-of-two. These power-of-two operations are necessary to calculate
        the number of integer bits used in the quantizer, and the difference
        between using round and ceil trade off the quantizer's range and
        precision.
      relu_neg_slope: Float. Slope of the negative values in relu to enable the
        use of leaky relu. This parameter will only be used with the quantizer
        type quantized_relu. Set to 0.0 to use normal relu.
      relu_upper_bound: Float. The upper bound to use if the activation is set
        to relu. Set to None to not artificially set an upper bound. Pease note
        that this param is ignored if the activation is not quantized_relu
      **kwargs: Args passed to the Layer class.
    """
    super(QAdaptiveActivation, self).__init__(**kwargs)

    self.total_bits = total_bits
    self.symmetric = symmetric
    self.is_estimating_step_count = False  # If the layer should estimate its
                                           # own step count by incrementing it
                                           # every call.
    if isinstance(current_step, tf.Variable):
      self.step = current_step
    elif current_step is None:
      self.step = tf.Variable(-1, dtype=tf.int64)
      self.is_estimating_step_count = True
      print("[WARNING] QAdaptiveActivation is estimating it's own training "
            "step count, which may not always be the same as the true optimizer"
            " training step. To mitigate this, please set the current_step "
            "parameter when initializing QAdaptiveActivation", file=sys.stderr)
    else:
      self.step = tf.Variable(current_step, dtype=tf.int64)
      print("[WARNING] QAdaptiveActivation is disconnected from the optimizer "
            "current step, which may lead to incorrect training. If you wish to"
            " resume training, set this layer's self.step to the optimizer's "
            "tf.Variable current step", file=sys.stderr)
    self.quantization_delay = quantization_delay
    self.ema_freeze_delay = ema_freeze_delay
    self.will_ema_freeze = True if ema_freeze_delay else False
    self.ema_decay = ema_decay
    self.per_channel = per_channel
    self.po2_rounding = po2_rounding
    self.ema_min = None
    self.ema_max = None
    self.relu_neg_slope = relu_neg_slope
    self.relu_upper_bound = relu_upper_bound

    # Verify quantizer type is correct
    self.supported_quantizers = ["quantized_bits", "quantized_relu"]
    if activation not in self.supported_quantizers:
      raise ValueError(("Invalid activation {}. Activation quantizer may NOT "
                        "contain any parameters (they will be set automatically"
                        " by this layer), and only the quantizer types {} are "
                        "supported.").format(activation,
                                             self.supported_quantizers))

    # Get the quantizer associated with the activation
    try:
      self.quantizer = get_quantizer(activation)
    except KeyError:
      raise ValueError("Invalid activation '{}'".format(activation))

    # Check that the quantizer is supported
    if self.quantizer.__class__.__name__ not in self.supported_quantizers:
      raise ValueError("Unsupported activation quantizer '{}'".format(
          self.quantizer.__class__.__name__))

    # Set keep_negative
    if self.quantizer.__class__.__name__ == "quantized_relu":
      self.quantizer.is_quantized_clip = False  # Use relu_upper_bound instead
      if self.relu_upper_bound:
        self.quantizer.relu_upper_bound = self.relu_upper_bound
      self.quantizer.negative_slope = relu_neg_slope
      self.keep_negative = relu_neg_slope != 0.0
      self.quantizer.is_quantized_clip = False  # Use normal relu when qnoise=0
    elif self.quantizer.__class__.__name__ == "quantized_bits":
      self.keep_negative = True
      self.quantizer.keep_negative = True

    # If not using quantization delay, then print warning
    if self.quantization_delay < 1:
      print("[WARNING] If QAdaptiveActivation has the quantization_delay set "
            "to 0, then the moving averages will be heavily biased towards the "
            "initial quantizer configuration, which will likely prevent the "
            "model from converging. Consider a larger quantization_delay.",
            file=sys.stderr)

    self.activation = self.quantizer  # self.activation is used by QTools

  def build(self, input_shape):
    if self.will_ema_freeze:
      self.ema_freeze_delay = tf.constant(self.ema_freeze_delay, dtype=tf.int64)

    self.ema_decay = tf.constant(self.ema_decay, dtype=tf.float32)
    self.is_estimating_step_count = tf.constant(self.is_estimating_step_count,
                                                dtype=tf.bool)

    # Calculate the number of channels
    channel_index = -1 if K.image_data_format() == "channels_last" else 1
    if self.per_channel:
      input_shape_list = list(input_shape) if isinstance(
          input_shape, tuple) else input_shape.as_list()
      num_channels = tf.constant(input_shape_list[channel_index],
                                 shape=(1), dtype=tf.int64)
    else:
      num_channels = tf.constant(1, shape=(1), dtype=tf.int64)

    # Initialize the moving mins and max
    if self.ema_min is None or self.ema_max is None:
      self.ema_min = tf.Variable(tf.zeros(num_channels), name="ema_min",
                                 trainable=False)
      self.ema_max = tf.Variable(tf.zeros(num_channels), name="ema_max",
                                 trainable=False)

    # Determine the parameters for the quantizer
    self.quantizer.bits = self.total_bits

    # Set up the initial integer bits and quantizer params
    self.quantizer.integer = tf.Variable(tf.zeros(num_channels,
                                                  dtype=tf.int32),
                                         name="quantizer_integer_bits",
                                         trainable=False)
    integer_bits = _get_integer_bits(min_value=self.ema_min,
                                     max_value=self.ema_max,
                                     bits=self.total_bits,
                                     symmetric=self.symmetric,
                                     keep_negative=self.keep_negative,
                                     is_clipping=self.po2_rounding)
    self.quantizer.integer.assign(integer_bits)
    self.quantizer.alpha = 1.0  # Setting alpha to 1.0 allows the integer bits
                                # to serve as the scale
    self.quantizer.symmetric = self.symmetric
    self.quantization_delay = tf.constant(self.quantization_delay,
                                          dtype=tf.int64)

  def call(self, inputs, training=False):
    x = inputs
    training = training and self.trainable
    self.will_ema_freeze = self.will_ema_freeze and self.trainable

    # Update the step count if the optimizer step count is unknown
    self.step.assign_add(K.switch(
        tf.math.logical_and(self.is_estimating_step_count, training),
        tf.constant(1, tf.int64), tf.constant(0, tf.int64)))

    # Perform the quantization
    if training:
      # Calculate the qnoise, a scalar from 0 to 1 that represents the level of
      # quantization noise to use. At training start, we want no quantization,
      # so qnoise_factor = 0.0. After quantization_delay steps, we want normal
      # quantization, so qnoise_factor = 1.0.
      qnoise_factor = K.switch(
          tf.greater_equal(self.step, self.quantization_delay),
          lambda: tf.constant(1.0), lambda: tf.constant(0.0))
      self.quantizer.update_qnoise_factor(qnoise_factor)
      qx = self.quantizer(x)

    else:  # If not training, we always want to use full quantization
      self.quantizer.update_qnoise_factor(tf.constant(1.0))
      qx = self.quantizer(x)

    # Calculate the axis along where to find the min and max EMAs
    len_axis = len(x.shape)
    if len_axis > 1:
      if self.per_channel:
        if K.image_data_format() == "channels_last":
          axis = list(range(len_axis - 1))
        else:
          axis = list(range(1, len_axis))
      else:
        axis = list(range(len_axis))
    else:
      axis = [0]

    # Determine if freezing the EMA
    is_ema_training = tf.constant(training, dtype=tf.bool)
    if self.will_ema_freeze:
      is_ema_training = tf.cond(
          tf.greater(self.step, self.ema_freeze_delay),
          lambda: tf.constant(False), lambda: tf.constant(True))

    def update_branch():
      """ Update the moving average when is_ema_training is True."""

      # Set the qnoise factor to 0 to update the EMA using the unquantized input
      prev_qnoise_factor = tf.identity(self.quantizer.qnoise_factor)
      self.quantizer.update_qnoise_factor(tf.constant(0.0))

      # Update the EMA
      act_x = self.quantizer(x)  # act_x is the input after the activation
                                 # function, but before the quantizer. This is
                                 # done by using a qnoise_factor of 0
      new_min = tf.squeeze(K.min(act_x, axis=axis, keepdims=True))
      K.moving_average_update(self.ema_min, new_min, self.ema_decay)
      new_max = tf.squeeze(K.max(act_x, axis=axis, keepdims=True))
      K.moving_average_update(self.ema_max, new_max, self.ema_decay)

      # Reset the qnoise factor to the previous value
      self.quantizer.update_qnoise_factor(prev_qnoise_factor)

    # Update the moving average when is_ema_training is True
    tf_utils.smart_cond(
        is_ema_training, true_fn=update_branch, false_fn=lambda: None)

    # Set the integer bits for the quantizer
    integer_bits = _get_integer_bits(
        min_value=self.ema_min,
        max_value=self.ema_max,
        bits=self.total_bits,
        symmetric=self.symmetric,
        keep_negative=self.keep_negative,
        is_clipping=self.po2_rounding)
    self.quantizer.integer.assign(integer_bits)

    return qx

  # Override get_weights since we do not want ema_min or ema_max to be public
  def get_weights(self):
    return []

  # Override set_weights since we do not want ema_min or ema_max to be public
  def set_weights(self, weights):
    return

  def get_config(self):
    config = {
        "activation": self.quantizer.__class__.__name__,
        "total_bits": self.total_bits,
        "current_step": self.step.numpy(),
        "symmetric": self.symmetric,
        "quantization_delay": np.array(self.quantization_delay),
        "ema_freeze_delay": np.array(self.ema_freeze_delay),
        "ema_decay": np.array(self.ema_decay),
        "per_channel": self.per_channel,
        "po2_rounding": self.po2_rounding,
        "relu_neg_slope": self.relu_neg_slope
    }
    base_config = super(QAdaptiveActivation, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    self.quantizer.integer_bits = np.array(self.quantizer)
    return str(self.quantizer)

  def compute_output_shape(self, input_shape):
    return input_shape

  def get_prunable_weights(self):
    return [] class QAveragePooling2D(AveragePooling2D):
  """Computes the quantized version of AveragePooling2D."""

  def __init__(self, pool_size=(2, 2),
               strides=None,
               padding="valid",
               data_format=None,
               average_quantizer=None,
               activation=None,
               **kwargs):

    self.average_quantizer = average_quantizer
    self.average_quantizer_internal = get_quantizer(self.average_quantizer)
    self.quantizers = [self.average_quantizer_internal]

    if activation is not None:
      self.activation = get_quantizer(activation)
    else:
      self.activation = activation

    super(QAveragePooling2D, self).__init__(
        pool_size=pool_size, strides=strides, padding=padding,
        data_format=data_format, **kwargs)

  def call(self, inputs):
    """Performs quantized AveragePooling followed by QActivation.

    Since there is no specific parameter for averaging op, we couldn't apply
    averaging quantizer to the averaging op. We have two options:
    1. we perform our own average as sum first then multiply with the
       inversion
       of the division factor: sum(x) * quantize(1/pool_area)
    2. first, we call keras version of averaging first: y1 = keras_average(x)
       then multiply it with pool_size^2: y2 = y1 * pool_area
       Last, y3 = y2 * quantize(1/ pool_area)
    Our numerical anaysis suggests negligible error between 1 and 2. Therefore
    we use option #2 here for the simplicity of implementation.

    """

    x = super(QAveragePooling2D, self).call(inputs)

    if self.average_quantizer:
      if isinstance(self.pool_size, int):
        pool_area = self.pool_size * self.pool_size
      else:
        pool_area = np.prod(self.pool_size)

      # Revertes the division results.
      x = x * pool_area

      # Quantizes the multiplication factor.
      mult_factor = 1.0 / pool_area

      q_mult_factor = self.average_quantizer_internal(mult_factor)
      q_mult_factor = K.cast_to_floatx(q_mult_factor)
      x = x * q_mult_factor

    if self.activation is not None:
      return self.activation(x)
    return x

  def get_config(self):
    config = {
        "average_quantizer":
            constraints.serialize(self.average_quantizer_internal),
        "activation": constraints.serialize(self.activation),
    }
    base_config = super(QAveragePooling2D, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return {
        "average_quantizer":
            str(self.average_quantizer_internal),
        "activation":
            str(self.activation)
    }

  def get_quantizers(self):
    return self.quantizers class QDense(Dense, PrunableLayer):
  """Implements a quantized Dense layer."""

  # Most of these parameters follow the implementation of Dense in
  # Keras, with the exception of kernel_range, bias_range,
  # kernel_quantizer, bias_quantizer, and kernel_initializer.
  #
  # kernel_quantizer: quantizer function/class for kernel
  # bias_quantizer: quantizer function/class for bias
  # kernel_range/bias_ranger: for quantizer functions whose values
  #   can go over [-1,+1], these values are used to set the clipping
  #   value of kernels and biases, respectively, instead of using the
  #   constraints specified by the user.
  #
  # we refer the reader to the documentation of Dense in Keras for the
  # other parameters.

  def __init__(self,
               units,
               activation=None,
               use_bias=True,
               kernel_initializer="he_normal",
               bias_initializer="zeros",
               kernel_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               bias_quantizer=None,
               kernel_range=None,
               bias_range=None,
               **kwargs):

    if kernel_range is not None:
      warnings.warn("kernel_range is deprecated in QDense layer.")

    if bias_range is not None:
      warnings.warn("bias_range is deprecated in QDense layer.")

    self.kernel_range = kernel_range
    self.bias_range = bias_range

    self.kernel_quantizer = kernel_quantizer
    self.bias_quantizer = bias_quantizer

    self.kernel_quantizer_internal = get_quantizer(self.kernel_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)

    # optimize parameter set to "auto" scaling mode if possible
    if hasattr(self.kernel_quantizer_internal, "_set_trainable_parameter"):
      self.kernel_quantizer_internal._set_trainable_parameter()

    self.quantizers = [
        self.kernel_quantizer_internal, self.bias_quantizer_internal
    ]

    kernel_constraint, kernel_initializer = (
        get_auto_range_constraint_initializer(self.kernel_quantizer_internal,
                                              kernel_constraint,
                                              kernel_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))
    if activation is not None:
      activation = get_quantizer(activation)

    super(QDense, self).__init__(
        units=units,
        activation=activation,
        use_bias=use_bias,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        kernel_constraint=kernel_constraint,
        bias_constraint=bias_constraint,
        **kwargs)

  def call(self, inputs):
    if self.kernel_quantizer:
      quantized_kernel = self.kernel_quantizer_internal(self.kernel)
    else:
      quantized_kernel = self.kernel
    output = tf.keras.backend.dot(inputs, quantized_kernel)
    if self.use_bias:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias
      output = tf.keras.backend.bias_add(output, quantized_bias,
                                         data_format="channels_last")
    if self.activation is not None:
      output = self.activation(output)
    return output

  def compute_output_shape(self, input_shape):
    assert input_shape and len(input_shape) >= 2
    assert input_shape[-1]
    output_shape = list(input_shape)
    output_shape[-1] = self.units
    return tuple(output_shape)

  def get_config(self):
    config = {
        "units": self.units,
        "activation": activations.serialize(self.activation),
        "use_bias": self.use_bias,
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "kernel_initializer":
            initializers.serialize(self.kernel_initializer),
        "bias_initializer":
            initializers.serialize(self.bias_initializer),
        "kernel_regularizer":
            regularizers.serialize(self.kernel_regularizer),
        "bias_regularizer":
            regularizers.serialize(self.bias_regularizer),
        "activity_regularizer":
            regularizers.serialize(self.activity_regularizer),
        "kernel_constraint":
            constraints.serialize(self.kernel_constraint),
        "bias_constraint":
            constraints.serialize(self.bias_constraint),
        "kernel_range": self.kernel_range,
        "bias_range": self.bias_range
    }
    base_config = super(QDense, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return {
        "kernel_quantizer":
            str(self.kernel_quantizer_internal),
        "bias_quantizer":
            str(self.bias_quantizer_internal),
        "activation":
            str(self.activation),
        "units" : str(self.units)
    }

  def get_quantizers(self):
    return self.quantizers

  def get_prunable_weights(self):
    return [self.kernel] class QInitializer(Initializer):
  """Wraps around Keras initializer to provide a fanin scaling factor."""

  def __init__(self, initializer, use_scale, quantizer):
    self.initializer = initializer
    self.use_scale = use_scale
    self.quantizer = quantizer

    try:
      self.is_po2 = "po2" in quantizer.__class__.__name__
    except:
      self.is_po2 = False

  def __call__(self, shape, dtype=None):
    x = self.initializer(shape, dtype)

    max_x = np.max(abs(x))
    std_x = np.std(x)
    delta = self.quantizer.max() * 2**-self.quantizer.bits

    # delta is the minimum resolution of the number system.
    # we want to make sure we have enough values.
    if delta > std_x and hasattr(self.initializer, "scale"):
      q = self.quantizer(x)
      max_q = np.max(abs(q))
      scale = 1.0
      if max_q == 0.0:
        xx = np.mean(x * x)
        scale = self.quantizer.max() / np.sqrt(xx)
      else:
        qx = np.sum(q * x)
        qq = np.sum(q * q)

        scale = qq / qx

      self.initializer.scale *= max(scale, 1)
      x = self.initializer(shape, dtype)

    return np.clip(x, -self.quantizer.max(), self.quantizer.max())

  def get_config(self):
    return {
        "initializer": self.initializer,
        "use_scale": self.use_scale,
        "quantizer": self.quantizer,
    }

  @classmethod
  def from_config(cls, config):
    config = {
      'initializer' : get_initializer(config['initializer']),
      'use_scale'   : config['use_scale'],
      'quantizer'   : get_quantizer(config['quantizer'])}
    return cls(**config) class QConv1D(Conv1D, PrunableLayer):
  """1D convolution layer (e.g. spatial convolution over images)."""

  # most of these parameters follow the implementation of Conv1D in Keras,
  # with the exception of kernel_range, bias_range, kernel_quantizer
  # and bias_quantizer, and kernel_initializer.
  #
  # kernel_quantizer: quantizer function/class for kernel
  # bias_quantizer: quantizer function/class for bias
  # kernel_range/bias_ranger: for quantizer functions whose values
  #   can go over [-1,+1], these values are used to set the clipping
  #   value of kernels and biases, respectively, instead of using the
  #   constraints specified by the user.
  #
  # we refer the reader to the documentation of Conv1D in Keras for the
  # other parameters.
  #

  def __init__(self,
               filters,
               kernel_size,
               strides=1,
               padding="valid",
               dilation_rate=1,
               activation=None,
               use_bias=True,
               kernel_initializer="he_normal",
               bias_initializer="zeros",
               kernel_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               bias_quantizer=None,
               kernel_range=None,
               bias_range=None,
               **kwargs):

    if kernel_range is not None:
      warnings.warn("kernel_range is deprecated in QConv1D layer.")

    if bias_range is not None:
      warnings.warn("bias_range is deprecated in QConv1D layer.")

    self.kernel_range = kernel_range
    self.bias_range = bias_range

    self.kernel_quantizer = kernel_quantizer
    self.bias_quantizer = bias_quantizer

    self.kernel_quantizer_internal = get_quantizer(self.kernel_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)

    # optimize parameter set to "auto" scaling mode if possible
    if hasattr(self.kernel_quantizer_internal, "_set_trainable_parameter"):
      self.kernel_quantizer_internal._set_trainable_parameter()

    self.quantizers = [
        self.kernel_quantizer_internal, self.bias_quantizer_internal
    ]

    kernel_constraint, kernel_initializer = (
        get_auto_range_constraint_initializer(self.kernel_quantizer_internal,
                                              kernel_constraint,
                                              kernel_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))
    if activation is not None:
      activation = get_quantizer(activation)

    super(QConv1D, self).__init__(
        filters=filters,
        kernel_size=kernel_size,
        strides=strides,
        padding=padding,
        dilation_rate=dilation_rate,
        activation=activation,
        use_bias=use_bias,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        kernel_constraint=kernel_constraint,
        bias_constraint=bias_constraint,
        **kwargs)

  def call(self, inputs):
    if self.kernel_quantizer:
      quantized_kernel = self.kernel_quantizer_internal(self.kernel)
    else:
      quantized_kernel = self.kernel

    outputs = tf.keras.backend.conv1d(
        inputs,
        quantized_kernel,
        strides=self.strides[0],
        padding=self.padding,
        data_format=self.data_format,
        dilation_rate=self.dilation_rate[0])

    if self.use_bias:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias

      outputs = tf.keras.backend.bias_add(
          outputs, quantized_bias, data_format=self.data_format)

    if self.activation is not None:
      return self.activation(outputs)
    return outputs

  def get_config(self):
    config = {
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "kernel_range": self.kernel_range,
        "bias_range": self.bias_range
    }
    base_config = super(QConv1D, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return {
        "kernel_quantizer":
            str(self.kernel_quantizer_internal),
        "bias_quantizer":
            str(self.bias_quantizer_internal),
        "activation":
            str(self.activation),
        "filters" : str(self.filters)
    }

  def get_quantizers(self):
    return self.quantizers

  def get_prunable_weights(self):
    return [self.kernel] class QConv2D(Conv2D, PrunableLayer):
  """2D convolution layer (e.g. spatial convolution over images)."""

  # most of these parameters follow the implementation of Conv2D in Keras,
  # with the exception of kernel_range, bias_range, kernel_quantizer
  # and bias_quantizer, and kernel_initializer.
  #
  # kernel_quantizer: quantizer function/class for kernel
  # bias_quantizer: quantizer function/class for bias
  # kernel_range/bias_ranger: for quantizer functions whose values
  #   can go over [-1,+1], these values are used to set the clipping
  #   value of kernels and biases, respectively, instead of using the
  #   constraints specified by the user.
  #
  # we refer the reader to the documentation of Conv2D in Keras for the
  # other parameters.
  #

  def __init__(self,
               filters,
               kernel_size,
               strides=(1, 1),
               padding="valid",
               data_format="channels_last",
               dilation_rate=(1, 1),
               activation=None,
               use_bias=True,
               kernel_initializer="he_normal",
               bias_initializer="zeros",
               kernel_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               bias_constraint=None,
               kernel_range=None,
               bias_range=None,
               kernel_quantizer=None,
               bias_quantizer=None,
               **kwargs):

    if kernel_range is not None:
      warnings.warn("kernel_range is deprecated in QConv2D layer.")

    if bias_range is not None:
      warnings.warn("bias_range is deprecated in QConv2D layer.")

    self.kernel_range = kernel_range
    self.bias_range = bias_range

    self.kernel_quantizer = kernel_quantizer
    self.bias_quantizer = bias_quantizer

    self.kernel_quantizer_internal = get_quantizer(self.kernel_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)

    # optimize parameter set to "auto" scaling mode if possible
    if hasattr(self.kernel_quantizer_internal, "_set_trainable_parameter"):
      self.kernel_quantizer_internal._set_trainable_parameter()

    self.quantizers = [
        self.kernel_quantizer_internal, self.bias_quantizer_internal
    ]

    kernel_constraint, kernel_initializer = (
        get_auto_range_constraint_initializer(self.kernel_quantizer_internal,
                                              kernel_constraint,
                                              kernel_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))

    if activation is not None:
      activation = get_quantizer(activation)

    super(QConv2D, self).__init__(
        filters=filters,
        kernel_size=kernel_size,
        strides=strides,
        padding=padding,
        data_format=data_format,
        dilation_rate=dilation_rate,
        activation=activation,
        use_bias=use_bias,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        kernel_constraint=kernel_constraint,
        bias_constraint=bias_constraint,
        **kwargs)

  def call(self, inputs):
    if self.kernel_quantizer:
      quantized_kernel = self.kernel_quantizer_internal(self.kernel)
    else:
      quantized_kernel = self.kernel

    outputs = tf.keras.backend.conv2d(
        inputs,
        quantized_kernel,
        strides=self.strides,
        padding=self.padding,
        data_format=self.data_format,
        dilation_rate=self.dilation_rate)

    if self.use_bias:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias

      outputs = tf.keras.backend.bias_add(
          outputs, quantized_bias, data_format=self.data_format)

    if self.activation is not None:
      return self.activation(outputs)
    return outputs

  def get_config(self):
    config = {
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "kernel_range": self.kernel_range,
        "bias_range": self.bias_range
    }
    base_config = super(QConv2D, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return {
        "kernel_quantizer":
            str(self.kernel_quantizer_internal),
        "bias_quantizer":
            str(self.bias_quantizer_internal),
        "activation":
            str(self.activation),
        "filters" : str(self.filters)
    }

  def get_quantizers(self):
    return self.quantizers

  def get_prunable_weights(self):
    return [self.kernel] class QConv2DTranspose(Conv2DTranspose, PrunableLayer):
  """2D convolution layer (e.g. spatial convolution over images)."""

  # most of these parameters follow the implementation of Conv2DTranspose
  # in Keras, with the exception of kernel_quantizer and bias_quantizer
  # and kernel_initializer.
  #
  # kernel_quantizer: quantizer function/class for kernel
  # bias_quantizer: quantizer function/class for bias
  #
  # we refer the reader to the documentation of Conv2DTranspose in Keras for
  # the other parameters.
  #

  def __init__(self,
               filters,
               kernel_size,
               strides=(1, 1),
               padding='valid',
               output_padding=None,
               data_format=None,
               dilation_rate=(1, 1),
               activation=None,
               use_bias=True,
               kernel_initializer='glorot_uniform',
               bias_initializer='zeros',
               kernel_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               bias_quantizer=None,
               **kwargs):

    self.kernel_quantizer = kernel_quantizer
    self.bias_quantizer = bias_quantizer

    self.kernel_quantizer_internal = get_quantizer(self.kernel_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)

    # optimize parameter set to "auto" scaling mode if possible
    if hasattr(self.kernel_quantizer_internal, "_set_trainable_parameter"):
      self.kernel_quantizer_internal._set_trainable_parameter()

    self.quantizers = [
        self.kernel_quantizer_internal, self.bias_quantizer_internal
    ]

    kernel_constraint, kernel_initializer = (
        get_auto_range_constraint_initializer(self.kernel_quantizer_internal,
                                              kernel_constraint,
                                              kernel_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))

    if activation is not None:
      activation = get_quantizer(activation)

    super(QConv2DTranspose, self).__init__(
        filters=filters,
        kernel_size=kernel_size,
        strides=strides,
        padding=padding,
        output_padding=None,
        data_format=data_format,
        dilation_rate=dilation_rate,
        activation=activation,
        use_bias=use_bias,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        kernel_constraint=kernel_constraint,
        bias_constraint=bias_constraint,
        **kwargs)

  def call(self, inputs):
    inputs_shape = array_ops.shape(inputs)
    batch_size = inputs_shape[0]
    if self.data_format == 'channels_first':
      h_axis, w_axis = 2, 3
    else:
      h_axis, w_axis = 1, 2

    height, width = inputs_shape[h_axis], inputs_shape[w_axis]
    kernel_h, kernel_w = self.kernel_size
    stride_h, stride_w = self.strides

    if self.output_padding is None:
      out_pad_h = out_pad_w = None
    else:
      out_pad_h, out_pad_w = self.output_padding

    # Infer the dynamic output shape:
    out_height = conv_utils.deconv_output_length(height,
                                                 kernel_h,
                                                 padding=self.padding,
                                                 output_padding=out_pad_h,
                                                 stride=stride_h,
                                                 dilation=self.dilation_rate[0])
    out_width = conv_utils.deconv_output_length(width,
                                                kernel_w,
                                                padding=self.padding,
                                                output_padding=out_pad_w,
                                                stride=stride_w,
                                                dilation=self.dilation_rate[1])
    if self.data_format == 'channels_first':
      output_shape = (batch_size, self.filters, out_height, out_width)
    else:
      output_shape = (batch_size, out_height, out_width, self.filters)

    if self.kernel_quantizer:
      quantized_kernel = self.kernel_quantizer_internal(self.kernel)
    else:
      quantized_kernel = self.kernel

    output_shape_tensor = array_ops.stack(output_shape)
    outputs = tf.keras.backend.conv2d_transpose(
        inputs,
        quantized_kernel,
        output_shape_tensor,
        strides=self.strides,
        padding=self.padding,
        data_format=self.data_format,
        dilation_rate=self.dilation_rate)

    if not context.executing_eagerly():
      # Infer the static output shape:
      out_shape = self.compute_output_shape(inputs.shape)
      outputs.set_shape(out_shape)

    if self.use_bias:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias

      outputs = tf.keras.backend.bias_add(
          outputs,
          quantized_bias,
          data_format=self.data_format)

    if self.activation is not None:
      return self.activation(outputs)
    return outputs

  def get_config(self):
    config = {
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal)
    }
    base_config = super(QConv2DTranspose, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantizers(self):
    return self.quantizers

  def get_prunable_weights(self):
    return [self.kernel] class QSimpleRNN(RNN, PrunableLayer):
  """
  Class for the QSimpleRNN layer.

  Most of these parameters follow the implementation of SimpleRNN in
  Keras, with the exception of kernel_quantizer, recurrent_quantizer,
  bias_quantizer and state_quantizer.


  kernel_quantizer: quantizer function/class for kernel
  recurrent_quantizer: quantizer function/class for recurrent kernel
  bias_quantizer: quantizer function/class for bias
  state_quantizer: quantizer function/class for states


  We refer the reader to the documentation of SimpleRNN in Keras for the
  other parameters.

  """

  def __init__(self,
               units,
               activation='quantized_tanh',
               use_bias=True,
               kernel_initializer='glorot_uniform',
               recurrent_initializer='orthogonal',
               bias_initializer='zeros',
               kernel_regularizer=None,
               recurrent_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               recurrent_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               recurrent_quantizer=None,
               bias_quantizer=None,
               state_quantizer=None,
               dropout=0.,
               recurrent_dropout=0.,
               return_sequences=False,
               return_state=False,
               go_backwards=False,
               stateful=False,
               unroll=False,
               **kwargs):

    if 'enable_caching_device' in kwargs:
      cell_kwargs = {'enable_caching_device':
                     kwargs.pop('enable_caching_device')}
    else:
      cell_kwargs = {}

    cell = QSimpleRNNCell(
        units,
        activation=activation,
        use_bias=use_bias,
        kernel_initializer=kernel_initializer,
        recurrent_initializer=recurrent_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        recurrent_regularizer=recurrent_regularizer,
        bias_regularizer=bias_regularizer,
        kernel_constraint=kernel_constraint,
        recurrent_constraint=recurrent_constraint,
        bias_constraint=bias_constraint,
        kernel_quantizer=kernel_quantizer,
        recurrent_quantizer=recurrent_quantizer,
        bias_quantizer=bias_quantizer,
        state_quantizer=state_quantizer,
        dropout=dropout,
        recurrent_dropout=recurrent_dropout,
        dtype=kwargs.get('dtype'),
        trainable=kwargs.get('trainable', True),
        **cell_kwargs)

    super(QSimpleRNN, self).__init__(
        cell,
        return_sequences=return_sequences,
        return_state=return_state,
        go_backwards=go_backwards,
        stateful=stateful,
        unroll=unroll,
        **kwargs)
    self.activity_regularizer = regularizers.get(activity_regularizer)
    self.input_spec = [tf.keras.layers.InputSpec(ndim=3)]

  def call(self, inputs, mask=None, training=None, initial_state=None):
    self._maybe_reset_cell_dropout_mask(self.cell)
    return super(QSimpleRNN, self).call(
        inputs, mask=mask, training=training, initial_state=initial_state)

  def get_quantizers(self):
    return self.cell.quantizers

  def get_prunable_weights(self):
    return [self.cell.kernel, self.cell.recurrent_kernel]

  @property
  def units(self):
    return self.cell.units

  @property
  def activation(self):
    return self.cell.activation

  @property
  def use_bias(self):
    return self.cell.use_bias

  @property
  def kernel_initializer(self):
    return self.cell.kernel_initializer

  @property
  def recurrent_initializer(self):
    return self.cell.recurrent_initializer

  @property
  def bias_initializer(self):
    return self.cell.bias_initializer

  @property
  def kernel_regularizer(self):
    return self.cell.kernel_regularizer

  @property
  def recurrent_regularizer(self):
    return self.cell.recurrent_regularizer

  @property
  def bias_regularizer(self):
    return self.cell.bias_regularizer

  @property
  def kernel_constraint(self):
    return self.cell.kernel_constraint

  @property
  def recurrent_constraint(self):
    return self.cell.recurrent_constraint

  @property
  def bias_constraint(self):
    return self.cell.bias_constraint

  @property
  def kernel_quantizer_internal(self):
    return self.cell.kernel_quantizer_internal

  @property
  def recurrent_quantizer_internal(self):
    return self.cell.recurrent_quantizer_internal

  @property
  def bias_quantizer_internal(self):
    return self.cell.bias_quantizer_internal

  @property
  def state_quantizer_internal(self):
    return self.cell.state_quantizer_internal

  @property
  def kernel_quantizer(self):
    return self.cell.kernel_quantizer

  @property
  def recurrent_quantizer(self):
    return self.cell.recurrent_quantizer

  @property
  def bias_quantizer(self):
    return self.cell.bias_quantizer

  @property
  def state_quantizer(self):
    return self.cell.state_quantizer

  @property
  def dropout(self):
    return self.cell.dropout

  @property
  def recurrent_dropout(self):
    return self.cell.recurrent_dropout

  def get_config(self):
    config = {
        'units':
            self.units,
        'activation':
            activations.serialize(self.activation),
        'use_bias':
            self.use_bias,
        'kernel_initializer':
            initializers.serialize(self.kernel_initializer),
        'recurrent_initializer':
            initializers.serialize(self.recurrent_initializer),
        'bias_initializer':
            initializers.serialize(self.bias_initializer),
        'kernel_regularizer':
            regularizers.serialize(self.kernel_regularizer),
        'recurrent_regularizer':
            regularizers.serialize(self.recurrent_regularizer),
        'bias_regularizer':
            regularizers.serialize(self.bias_regularizer),
        'activity_regularizer':
            regularizers.serialize(self.activity_regularizer),
        'kernel_constraint':
            constraints.serialize(self.kernel_constraint),
        'recurrent_constraint':
            constraints.serialize(self.recurrent_constraint),
        'bias_constraint':
            constraints.serialize(self.bias_constraint),
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            constraints.serialize(self.recurrent_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "state_quantizer":
            constraints.serialize(self.state_quantizer_internal),
        'dropout':
            self.dropout,
        'recurrent_dropout':
            self.recurrent_dropout
    }
    base_config = super(QSimpleRNN, self).get_config()
    del base_config['cell']
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return {
        "kernel_quantizer":
            str(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            str(self.recurrent_quantizer_internal),
        "bias_quantizer":
            str(self.bias_quantizer_internal),
        "state_quantizer":
            str(self.state_quantizer_internal),
        "activation":
            str(self.activation)
    }

  @classmethod
  def from_config(cls, config):
    if 'implementation' in config:
      config.pop('implementation')
    return cls(**config) class QSimpleRNNCell(SimpleRNNCell):
  """
  Cell class for the QSimpleRNNCell layer.

  Most of these parameters follow the implementation of SimpleRNNCell in
  Keras, with the exception of kernel_quantizer, recurrent_quantizer,
  bias_quantizer, and state_quantizer.

  kernel_quantizer: quantizer function/class for kernel
  recurrent_quantizer: quantizer function/class for recurrent kernel
  bias_quantizer: quantizer function/class for bias
  state_quantizer: quantizer function/class for states

  We refer the reader to the documentation of SimpleRNNCell in Keras for the
  other parameters.

  """
  def __init__(self,
               units,
               activation='quantized_tanh',
               use_bias=True,
               kernel_initializer='glorot_uniform',
               recurrent_initializer='orthogonal',
               bias_initializer='zeros',
               kernel_regularizer=None,
               recurrent_regularizer=None,
               bias_regularizer=None,
               kernel_constraint=None,
               recurrent_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               recurrent_quantizer=None,
               bias_quantizer=None,
               state_quantizer=None,
               dropout=0.,
               recurrent_dropout=0.,
               **kwargs):

    self.kernel_quantizer = kernel_quantizer
    self.recurrent_quantizer = recurrent_quantizer
    self.bias_quantizer = bias_quantizer
    self.state_quantizer = state_quantizer

    self.kernel_quantizer_internal = get_quantizer(self.kernel_quantizer)
    self.recurrent_quantizer_internal = get_quantizer(self.recurrent_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)
    self.state_quantizer_internal = get_quantizer(self.state_quantizer)

    self.quantizers = [
        self.kernel_quantizer_internal, self.recurrent_quantizer_internal,
        self.bias_quantizer_internal, self.state_quantizer_internal
    ]

    if hasattr(self.kernel_quantizer_internal, "_set_trainable_parameter"):
      self.kernel_quantizer_internal._set_trainable_parameter()

    if hasattr(self.recurrent_quantizer_internal, "_set_trainable_parameter"):
      self.recurrent_quantizer_internal._set_trainable_parameter()

    kernel_constraint, kernel_initializer = (
        get_auto_range_constraint_initializer(self.kernel_quantizer_internal,
                                              kernel_constraint,
                                              kernel_initializer))

    recurrent_constraint, recurrent_initializer = (
        get_auto_range_constraint_initializer(self.recurrent_quantizer_internal,
                                              recurrent_constraint,
                                              recurrent_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))

    if activation is not None:
      activation = get_quantizer(activation)

    super(QSimpleRNNCell, self).__init__(
      units=units,
      activation=activation,
      use_bias=use_bias,
      kernel_initializer=kernel_initializer,
      recurrent_initializer=recurrent_initializer,
      bias_initializer=bias_initializer,
      kernel_regularizer=kernel_regularizer,
      recurrent_regularizer=recurrent_regularizer,
      bias_regularizer=bias_regularizer,
      kernel_constraint=kernel_constraint,
      recurrent_constraint=recurrent_constraint,
      bias_constraint=bias_constraint,
      dropout=dropout,
      recurrent_dropout=recurrent_dropout,
      **kwargs
    )

  def call(self, inputs, states, training=None):
    prev_output = states[0] if nest.is_sequence(states) else states

    dp_mask = self.get_dropout_mask_for_cell(inputs, training)
    rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(
        prev_output, training)

    if self.state_quantizer:
      quantized_prev_output = self.state_quantizer_internal(prev_output)
    else:
      quantized_prev_output = prev_output

    if self.kernel_quantizer:
      quantized_kernel = self.kernel_quantizer_internal(self.kernel)
    else:
      quantized_kernel = self.kernel

    if dp_mask is not None:
      h = K.dot(inputs * dp_mask, quantized_kernel)
    else:
      h = K.dot(inputs, quantized_kernel)

    if self.bias is not None:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias

      h = K.bias_add(h, quantized_bias)

    if rec_dp_mask is not None:
      quantized_prev_output = quantized_prev_output * rec_dp_mask

    if self.recurrent_quantizer:
      quantized_recurrent = self.recurrent_quantizer_internal(self.recurrent_kernel)
    else:
      quantized_recurrent = self.recurrent_kernel

    output = h + K.dot(quantized_prev_output, quantized_recurrent)

    if self.activation is not None:
      output = self.activation(output)
    return output, [output]

  def get_config(self):
    config = {
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            constraints.serialize(self.recurrent_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "state_quantizer":
            constraints.serialize(self.state_quantizer_internal)
    }
    base_config = super(QSimpleRNNCell, self).get_config()
    return dict(list(base_config.items()) + list(config.items())) class QLSTM(RNN, PrunableLayer):
  """
  Class for the QLSTM layer.

  Most of these parameters follow the implementation of LSTM in
  Keras, with the exception of kernel_quantizer, recurrent_quantizer,
  bias_quantizer, state_quantizer.


  kernel_quantizer: quantizer function/class for kernel
  recurrent_quantizer: quantizer function/class for recurrent kernel
  bias_quantizer: quantizer function/class for bias
  state_quantizer: quantizer function/class for states

  We refer the reader to the documentation of LSTM in Keras for the
  other parameters.

  """

  def __init__(self,
               units,
               activation='quantized_tanh',
               recurrent_activation='hard_sigmoid',
               use_bias=True,
               kernel_initializer='glorot_uniform',
               recurrent_initializer='orthogonal',
               bias_initializer='zeros',
               unit_forget_bias=True,
               kernel_regularizer=None,
               recurrent_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               recurrent_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               recurrent_quantizer=None,
               bias_quantizer=None,
               state_quantizer=None,
               dropout=0.,
               recurrent_dropout=0.,
               implementation=1,
               return_sequences=False,
               return_state=False,
               go_backwards=False,
               stateful=False,
               unroll=False,
               **kwargs):
    if implementation == 0:
      print('`implementation=0` has been deprecated, '
              'and now defaults to `implementation=1`.'
              'Please update your layer call.')

    if 'enable_caching_device' in kwargs:
      cell_kwargs = {'enable_caching_device':
                     kwargs.pop('enable_caching_device')}
    else:
      cell_kwargs = {}

    cell = QLSTMCell(
        units,
        activation=activation,
        recurrent_activation=recurrent_activation,
        use_bias=use_bias,
        kernel_initializer=kernel_initializer,
        recurrent_initializer=recurrent_initializer,
        unit_forget_bias=unit_forget_bias,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        recurrent_regularizer=recurrent_regularizer,
        bias_regularizer=bias_regularizer,
        kernel_constraint=kernel_constraint,
        recurrent_constraint=recurrent_constraint,
        bias_constraint=bias_constraint,
        kernel_quantizer=kernel_quantizer,
        recurrent_quantizer=recurrent_quantizer,
        bias_quantizer=bias_quantizer,
        state_quantizer=state_quantizer,
        dropout=dropout,
        recurrent_dropout=recurrent_dropout,
        implementation=implementation,
        dtype=kwargs.get('dtype'),
        trainable=kwargs.get('trainable', True),
        **cell_kwargs)

    super(QLSTM, self).__init__(
        cell,
        return_sequences=return_sequences,
        return_state=return_state,
        go_backwards=go_backwards,
        stateful=stateful,
        unroll=unroll,
        **kwargs)
    self.activity_regularizer = regularizers.get(activity_regularizer)
    self.input_spec = [tf.keras.layers.InputSpec(ndim=3)]

  def call(self, inputs, mask=None, training=None, initial_state=None):
    self._maybe_reset_cell_dropout_mask(self.cell)
    return super(QLSTM, self).call(
        inputs, mask=mask, training=training, initial_state=initial_state)

  def get_quantizers(self):
    return self.cell.quantizers

  def get_prunable_weights(self):
    return [self.cell.kernel, self.cell.recurrent_kernel]

  @property
  def units(self):
    return self.cell.units

  @property
  def activation(self):
    return self.cell.activation

  @property
  def recurrent_activation(self):
    return self.cell.recurrent_activation

  @property
  def use_bias(self):
    return self.cell.use_bias

  @property
  def kernel_initializer(self):
    return self.cell.kernel_initializer

  @property
  def recurrent_initializer(self):
    return self.cell.recurrent_initializer

  @property
  def bias_initializer(self):
    return self.cell.bias_initializer

  @property
  def unit_forget_bias(self):
    return self.cell.unit_forget_bias

  @property
  def kernel_regularizer(self):
    return self.cell.kernel_regularizer

  @property
  def recurrent_regularizer(self):
    return self.cell.recurrent_regularizer

  @property
  def bias_regularizer(self):
    return self.cell.bias_regularizer

  @property
  def kernel_constraint(self):
    return self.cell.kernel_constraint

  @property
  def recurrent_constraint(self):
    return self.cell.recurrent_constraint

  @property
  def bias_constraint(self):
    return self.cell.bias_constraint

  @property
  def kernel_quantizer_internal(self):
    return self.cell.kernel_quantizer_internal

  @property
  def recurrent_quantizer_internal(self):
    return self.cell.recurrent_quantizer_internal

  @property
  def bias_quantizer_internal(self):
    return self.cell.bias_quantizer_internal

  @property
  def state_quantizer_internal(self):
    return self.cell.state_quantizer_internal

  @property
  def kernel_quantizer(self):
    return self.cell.kernel_quantizer

  @property
  def recurrent_quantizer(self):
    return self.cell.recurrent_quantizer

  @property
  def bias_quantizer(self):
    return self.cell.bias_quantizer

  @property
  def state_quantizer(self):
    return self.cell.state_quantizer

  @property
  def dropout(self):
    return self.cell.dropout

  @property
  def recurrent_dropout(self):
    return self.cell.recurrent_dropout

  @property
  def implementation(self):
    return self.cell.implementation

  def get_config(self):
    config = {
        'units':
            self.units,
        'activation':
            activations.serialize(self.activation),
        'recurrent_activation':
            activations.serialize(self.recurrent_activation),
        'use_bias':
            self.use_bias,
        'kernel_initializer':
            initializers.serialize(self.kernel_initializer),
        'recurrent_initializer':
            initializers.serialize(self.recurrent_initializer),
        'bias_initializer':
            initializers.serialize(self.bias_initializer),
        'unit_forget_bias':
            self.unit_forget_bias,
        'kernel_regularizer':
            regularizers.serialize(self.kernel_regularizer),
        'recurrent_regularizer':
            regularizers.serialize(self.recurrent_regularizer),
        'bias_regularizer':
            regularizers.serialize(self.bias_regularizer),
        'activity_regularizer':
            regularizers.serialize(self.activity_regularizer),
        'kernel_constraint':
            constraints.serialize(self.kernel_constraint),
        'recurrent_constraint':
            constraints.serialize(self.recurrent_constraint),
        'bias_constraint':
            constraints.serialize(self.bias_constraint),
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            constraints.serialize(self.recurrent_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "state_quantizer":
            constraints.serialize(self.state_quantizer_internal),
        'dropout':
            self.dropout,
        'recurrent_dropout':
            self.recurrent_dropout,
        'implementation':
            self.implementation
    }
    base_config = super(QLSTM, self).get_config()
    del base_config['cell']
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return {
        "kernel_quantizer":
            str(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            str(self.recurrent_quantizer_internal),
        "bias_quantizer":
            str(self.bias_quantizer_internal),
        "state_quantizer":
            str(self.state_quantizer_internal),
        "activation":
            str(self.activation),
        "recurrent_activation":
            str(self.recurrent_activation),
    }

  @classmethod
  def from_config(cls, config):
    if 'implementation' in config and config['implementation'] == 0:
      config['implementation'] = 1
    return cls(**config) class QLSTMCell(LSTMCell):
  """
  Cell class for the QLSTMCell layer.

  Most of these parameters follow the implementation of LSTMCell in
  Keras, with the exception of kernel_quantizer, recurrent_quantizer,
  bias_quantizer, state_quantizer.


  kernel_quantizer: quantizer function/class for kernel
  recurrent_quantizer: quantizer function/class for recurrent kernel
  bias_quantizer: quantizer function/class for bias
  state_quantizer: quantizer function/class for states

  We refer the reader to the documentation of LSTMCell in Keras for the
  other parameters.

  """

  def __init__(self,
               units,
               activation='quantized_tanh',
               recurrent_activation='hard_sigmoid',
               use_bias=True,
               kernel_initializer='glorot_uniform',
               recurrent_initializer='orthogonal',
               bias_initializer='zeros',
               unit_forget_bias=True,
               kernel_regularizer=None,
               recurrent_regularizer=None,
               bias_regularizer=None,
               kernel_constraint=None,
               recurrent_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               recurrent_quantizer=None,
               bias_quantizer=None,
               state_quantizer=None,
               dropout=0.,
               recurrent_dropout=0.,
               implementation=1,
               **kwargs):
    self.kernel_quantizer = kernel_quantizer
    self.recurrent_quantizer = recurrent_quantizer
    self.bias_quantizer = bias_quantizer
    self.state_quantizer = state_quantizer

    self.kernel_quantizer_internal = get_quantizer(self.kernel_quantizer)
    self.recurrent_quantizer_internal = get_quantizer(self.recurrent_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)
    self.state_quantizer_internal = get_quantizer(self.state_quantizer)

    self.quantizers = [
      self.kernel_quantizer_internal,
      self.recurrent_quantizer_internal,
      self.bias_quantizer_internal,
      self.state_quantizer_internal,
    ]

    if hasattr(self.kernel_quantizer_internal, "_set_trainable_parameter"):
      self.kernel_quantizer_internal._set_trainable_parameter()

    if hasattr(self.recurrent_quantizer_internal, "_set_trainable_parameter"):
      self.recurrent_quantizer_internal._set_trainable_parameter()

    kernel_constraint, kernel_initializer = (
        get_auto_range_constraint_initializer(self.kernel_quantizer_internal,
                                              kernel_constraint,
                                              kernel_initializer))

    recurrent_constraint, recurrent_initializer = (
        get_auto_range_constraint_initializer(self.recurrent_quantizer_internal,
                                              recurrent_constraint,
                                              recurrent_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))

    if activation is not None:
      activation = get_quantizer(activation)

    if recurrent_activation is not None:
      recurrent_activation = get_quantizer(recurrent_activation)

    super(QLSTMCell, self).__init__(
      units=units,
      activation=activation,
      use_bias=use_bias,
      recurrent_activation=recurrent_activation,
      kernel_initializer=kernel_initializer,
      recurrent_initializer=recurrent_initializer,
      bias_initializer=bias_initializer,
      unit_forget_bias=True,
      kernel_regularizer=kernel_regularizer,
      recurrent_regularizer=recurrent_regularizer,
      bias_regularizer=bias_regularizer,
      kernel_constraint=kernel_constraint,
      recurrent_constraint=recurrent_constraint,
      bias_constraint=bias_constraint,
      dropout=dropout,
      recurrent_dropout=recurrent_dropout,
      implementation=implementation,
      **kwargs
    )

  def _compute_carry_and_output(self, x, h_tm1, c_tm1, quantized_recurrent):
    """Computes carry and output using split kernels."""
    x_i, x_f, x_c, x_o = x
    h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
    i = self.recurrent_activation(
        x_i + K.dot(h_tm1_i, quantized_recurrent[:, :self.units]))
    f = self.recurrent_activation(x_f + K.dot(
        h_tm1_f, quantized_recurrent[:, self.units:self.units * 2]))
    c = f * c_tm1 + i * self.activation(x_c + K.dot(
        h_tm1_c, quantized_recurrent[:, self.units * 2:self.units * 3]))
    o = self.recurrent_activation(
        x_o + K.dot(h_tm1_o, quantized_recurrent[:, self.units * 3:]))
    return c, o

  def _compute_carry_and_output_fused(self, z, c_tm1):
    """Computes carry and output using fused kernels."""
    z0, z1, z2, z3 = z
    i = self.recurrent_activation(z0)
    f = self.recurrent_activation(z1)
    c = f * c_tm1 + i * self.activation(z2)
    o = self.recurrent_activation(z3)
    return c, o

  def call(self, inputs, states, training=None):
    h_tm1_tmp = states[0]  # previous memory state
    c_tm1_tmp = states[1]  # previous carry state

    if self.state_quantizer:
      c_tm1 = self.state_quantizer_internal(c_tm1_tmp)
      h_tm1 = self.state_quantizer_internal(h_tm1_tmp)
    else:
      c_tm1 = c_tm1_tmp
      h_tm1 = h_tm1_tmp

    dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
    rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(
        h_tm1, training, count=4)

    if self.kernel_quantizer:
      quantized_kernel = self.kernel_quantizer_internal(self.kernel)
    else:
      quantized_kernel = self.kernel
    if self.recurrent_quantizer:
      quantized_recurrent = self.recurrent_quantizer_internal(self.recurrent_kernel)
    else:
      quantized_recurrent = self.recurrent_kernel
    if self.bias_quantizer:
      quantized_bias = self.bias_quantizer_internal(self.bias)
    else:
      quantized_bias = self.bias

    if self.implementation == 1:
      if 0 < self.dropout < 1.:
        inputs_i = inputs * dp_mask[0]
        inputs_f = inputs * dp_mask[1]
        inputs_c = inputs * dp_mask[2]
        inputs_o = inputs * dp_mask[3]
      else:
        inputs_i = inputs
        inputs_f = inputs
        inputs_c = inputs
        inputs_o = inputs
      k_i, k_f, k_c, k_o = array_ops.split(
          quantized_kernel, num_or_size_splits=4, axis=1)
      x_i = K.dot(inputs_i, k_i)
      x_f = K.dot(inputs_f, k_f)
      x_c = K.dot(inputs_c, k_c)
      x_o = K.dot(inputs_o, k_o)
      if self.use_bias:
        b_i, b_f, b_c, b_o = array_ops.split(
            quantized_bias, num_or_size_splits=4, axis=0)
        x_i = K.bias_add(x_i, b_i)
        x_f = K.bias_add(x_f, b_f)
        x_c = K.bias_add(x_c, b_c)
        x_o = K.bias_add(x_o, b_o)

      if 0 < self.recurrent_dropout < 1.:
        h_tm1_i = h_tm1 * rec_dp_mask[0]
        h_tm1_f = h_tm1 * rec_dp_mask[1]
        h_tm1_c = h_tm1 * rec_dp_mask[2]
        h_tm1_o = h_tm1 * rec_dp_mask[3]
      else:
        h_tm1_i = h_tm1
        h_tm1_f = h_tm1
        h_tm1_c = h_tm1
        h_tm1_o = h_tm1
      x = (x_i, x_f, x_c, x_o)
      h_tm1 = (h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o)
      c, o = self._compute_carry_and_output(x, h_tm1, c_tm1, quantized_recurrent)
    else:
      if 0. < self.dropout < 1.:
        inputs = inputs * dp_mask[0]
      z = K.dot(inputs, quantized_kernel)
      z += K.dot(h_tm1, quantized_recurrent)
      if self.use_bias:
        z = K.bias_add(z, quantized_bias)

      z = array_ops.split(z, num_or_size_splits=4, axis=1)
      c, o = self._compute_carry_and_output_fused(z, c_tm1)

    h = o * self.activation(c)
    return h, [h, c]

  def get_config(self):
    config = {
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            constraints.serialize(self.recurrent_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "state_quantizer":
            constraints.serialize(self.state_quantizer_internal)
    }
    base_config = super(QLSTMCell, self).get_config()
    return dict(list(base_config.items()) + list(config.items())) class QGRU(RNN, PrunableLayer):
  """
  Class for the QGRU layer.

  Most of these parameters follow the implementation of GRU in
  Keras, with the exception of kernel_quantizer, recurrent_quantizer,
  bias_quantizer and state_quantizer.


  kernel_quantizer: quantizer function/class for kernel
  recurrent_quantizer: quantizer function/class for recurrent kernel
  bias_quantizer: quantizer function/class for bias
  state_quantizer: quantizer function/class for states


  We refer the reader to the documentation of GRU in Keras for the
  other parameters.

  """

  def __init__(self,
               units,
               activation='quantized_tanh',
               recurrent_activation='hard_sigmoid',
               use_bias=True,
               kernel_initializer='glorot_uniform',
               recurrent_initializer='orthogonal',
               bias_initializer='zeros',
               kernel_regularizer=None,
               recurrent_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               recurrent_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               recurrent_quantizer=None,
               bias_quantizer=None,
               state_quantizer=None,
               dropout=0.,
               recurrent_dropout=0.,
               implementation=1,
               return_sequences=False,
               return_state=False,
               go_backwards=False,
               stateful=False,
               unroll=False,
               reset_after=False,
               **kwargs):
    if implementation == 0:
      print('`implementation=0` has been deprecated, '
              'and now defaults to `implementation=1`.'
              'Please update your layer call.')

    if 'enable_caching_device' in kwargs:
      cell_kwargs = {'enable_caching_device':
                     kwargs.pop('enable_caching_device')}
    else:
      cell_kwargs = {}

    cell = QGRUCell(
        units,
        activation=activation,
        recurrent_activation=recurrent_activation,
        use_bias=use_bias,
        kernel_initializer=kernel_initializer,
        recurrent_initializer=recurrent_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        recurrent_regularizer=recurrent_regularizer,
        bias_regularizer=bias_regularizer,
        kernel_constraint=kernel_constraint,
        recurrent_constraint=recurrent_constraint,
        bias_constraint=bias_constraint,
        kernel_quantizer=kernel_quantizer,
        recurrent_quantizer=recurrent_quantizer,
        bias_quantizer=bias_quantizer,
        state_quantizer=state_quantizer,
        dropout=dropout,
        recurrent_dropout=recurrent_dropout,
        implementation=implementation,
        reset_after=reset_after,
        dtype=kwargs.get('dtype'),
        trainable=kwargs.get('trainable', True),
        **cell_kwargs)

    super(QGRU, self).__init__(
        cell,
        return_sequences=return_sequences,
        return_state=return_state,
        go_backwards=go_backwards,
        stateful=stateful,
        unroll=unroll,
        **kwargs)
    self.activity_regularizer = regularizers.get(activity_regularizer)
    self.input_spec = [tf.keras.layers.InputSpec(ndim=3)]

  def call(self, inputs, mask=None, training=None, initial_state=None):
    self._maybe_reset_cell_dropout_mask(self.cell)
    return super(QGRU, self).call(
        inputs, mask=mask, training=training, initial_state=initial_state)

  def get_quantizers(self):
    return self.cell.quantizers

  def get_prunable_weights(self):
    return [self.cell.kernel, self.cell.recurrent_kernel]

  @property
  def units(self):
    return self.cell.units

  @property
  def activation(self):
    return self.cell.activation

  @property
  def recurrent_activation(self):
    return self.cell.recurrent_activation

  @property
  def use_bias(self):
    return self.cell.use_bias

  @property
  def kernel_initializer(self):
    return self.cell.kernel_initializer

  @property
  def recurrent_initializer(self):
    return self.cell.recurrent_initializer

  @property
  def bias_initializer(self):
    return self.cell.bias_initializer

  @property
  def kernel_regularizer(self):
    return self.cell.kernel_regularizer

  @property
  def recurrent_regularizer(self):
    return self.cell.recurrent_regularizer

  @property
  def bias_regularizer(self):
    return self.cell.bias_regularizer

  @property
  def kernel_constraint(self):
    return self.cell.kernel_constraint

  @property
  def recurrent_constraint(self):
    return self.cell.recurrent_constraint

  @property
  def bias_constraint(self):
    return self.cell.bias_constraint

  @property
  def kernel_quantizer_internal(self):
    return self.cell.kernel_quantizer_internal

  @property
  def recurrent_quantizer_internal(self):
    return self.cell.recurrent_quantizer_internal

  @property
  def bias_quantizer_internal(self):
    return self.cell.bias_quantizer_internal

  @property
  def state_quantizer_internal(self):
    return self.cell.state_quantizer_internal

  @property
  def kernel_quantizer(self):
    return self.cell.kernel_quantizer

  @property
  def recurrent_quantizer(self):
    return self.cell.recurrent_quantizer

  @property
  def bias_quantizer(self):
    return self.cell.bias_quantizer

  @property
  def state_quantizer(self):
    return self.cell.state_quantizer

  @property
  def dropout(self):
    return self.cell.dropout

  @property
  def recurrent_dropout(self):
    return self.cell.recurrent_dropout

  @property
  def implementation(self):
    return self.cell.implementation

  @property
  def reset_after(self):
    return self.cell.reset_after

  def get_config(self):
    config = {
        'units':
            self.units,
        'activation':
            activations.serialize(self.activation),
        'recurrent_activation':
            activations.serialize(self.recurrent_activation),
        'use_bias':
            self.use_bias,
        'kernel_initializer':
            initializers.serialize(self.kernel_initializer),
        'recurrent_initializer':
            initializers.serialize(self.recurrent_initializer),
        'bias_initializer':
            initializers.serialize(self.bias_initializer),
        'kernel_regularizer':
            regularizers.serialize(self.kernel_regularizer),
        'recurrent_regularizer':
            regularizers.serialize(self.recurrent_regularizer),
        'bias_regularizer':
            regularizers.serialize(self.bias_regularizer),
        'activity_regularizer':
            regularizers.serialize(self.activity_regularizer),
        'kernel_constraint':
            constraints.serialize(self.kernel_constraint),
        'recurrent_constraint':
            constraints.serialize(self.recurrent_constraint),
        'bias_constraint':
            constraints.serialize(self.bias_constraint),
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            constraints.serialize(self.recurrent_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "state_quantizer":
            constraints.serialize(self.state_quantizer_internal),
        'dropout':
            self.dropout,
        'recurrent_dropout':
            self.recurrent_dropout,
        'implementation':
            self.implementation,
        'reset_after':
            self.reset_after
    }
    base_config = super(QGRU, self).get_config()
    del base_config['cell']
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return {
        "kernel_quantizer":
            str(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            str(self.recurrent_quantizer_internal),
        "bias_quantizer":
            str(self.bias_quantizer_internal),
        "state_quantizer":
            str(self.state_quantizer_internal),
        "activation":
            str(self.activation),
        "recurrent_activation":
            str(self.recurrent_activation),
    }

  @classmethod
  def from_config(cls, config):
    if 'implementation' in config and config['implementation'] == 0:
      config['implementation'] = 1
    return cls(**config) class QGRUCell(GRUCell):
  """
  Cell class for the QGRUCell layer.

  Most of these parameters follow the implementation of GRUCell in
  Keras, with the exception of kernel_quantizer, recurrent_quantizer,
  bias_quantizer and state_quantizer.


  kernel_quantizer: quantizer function/class for kernel
  recurrent_quantizer: quantizer function/class for recurrent kernel
  bias_quantizer: quantizer function/class for bias
  state_quantizer: quantizer function/class for states


  We refer the reader to the documentation of GRUCell in Keras for the
  other parameters.

  """
  def __init__(self,
               units,
               activation='quantized_tanh',
               recurrent_activation='hard_sigmoid',
               use_bias=True,
               kernel_initializer='glorot_uniform',
               recurrent_initializer='orthogonal',
               bias_initializer='zeros',
               kernel_regularizer=None,
               recurrent_regularizer=None,
               bias_regularizer=None,
               kernel_constraint=None,
               recurrent_constraint=None,
               bias_constraint=None,
               kernel_quantizer=None,
               recurrent_quantizer=None,
               bias_quantizer=None,
               state_quantizer=None,
               dropout=0.,
               recurrent_dropout=0.,
               implementation=1,
               reset_after=False,
               **kwargs):

    self.kernel_quantizer = kernel_quantizer
    self.recurrent_quantizer = recurrent_quantizer
    self.bias_quantizer = bias_quantizer
    self.state_quantizer = state_quantizer

    self.kernel_quantizer_internal = get_quantizer(self.kernel_quantizer)
    self.recurrent_quantizer_internal = get_quantizer(self.recurrent_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)
    self.state_quantizer_internal = get_quantizer(self.state_quantizer)

    self.quantizers = [
      self.kernel_quantizer_internal,
      self.recurrent_quantizer_internal,
      self.bias_quantizer_internal,
      self.state_quantizer_internal
    ]

    if hasattr(self.kernel_quantizer_internal, "_set_trainable_parameter"):
      self.kernel_quantizer_internal._set_trainable_parameter()

    if hasattr(self.recurrent_quantizer_internal, "_set_trainable_parameter"):
      self.recurrent_quantizer_internal._set_trainable_parameter()

    kernel_constraint, kernel_initializer = (
        get_auto_range_constraint_initializer(self.kernel_quantizer_internal,
                                              kernel_constraint,
                                              kernel_initializer))

    recurrent_constraint, recurrent_initializer = (
        get_auto_range_constraint_initializer(self.recurrent_quantizer_internal,
                                              recurrent_constraint,
                                              recurrent_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))

    if activation is not None:
      activation = get_quantizer(activation)

    if recurrent_activation is not None:
      recurrent_activation = get_quantizer(recurrent_activation)

    super(QGRUCell, self).__init__(
      units=units,
      activation=activation,
      recurrent_activation=recurrent_activation,
      use_bias=use_bias,
      kernel_initializer=kernel_initializer,
      recurrent_initializer=recurrent_initializer,
      bias_initializer=bias_initializer,
      kernel_regularizer=kernel_regularizer,
      recurrent_regularizer=recurrent_regularizer,
      bias_regularizer=bias_regularizer,
      kernel_constraint=kernel_constraint,
      recurrent_constraint=recurrent_constraint,
      bias_constraint=bias_constraint,
      dropout=dropout,
      recurrent_dropout=recurrent_dropout,
      implementation=implementation,
      reset_after=reset_after,
      **kwargs)

  def call(self, inputs, states, training=None):
    # previous memory
    h_tm1_tmp = states[0] if nest.is_sequence(states) else states

    dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=3)
    rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(
        h_tm1_tmp, training, count=3)

    if self.state_quantizer:
      h_tm1 = self.state_quantizer_internal(h_tm1_tmp)
    else:
      h_tm1 = h_tm1_tmp

    if self.kernel_quantizer:
      quantized_kernel = self.kernel_quantizer_internal(self.kernel)
    else:
      quantized_kernel = self.kernel
    if self.recurrent_quantizer:
      quantized_recurrent = self.recurrent_quantizer_internal(self.recurrent_kernel)
    else:
      quantized_recurrent = self.kernel

    if self.use_bias:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias

      if not self.reset_after:
        input_bias, recurrent_bias = quantized_bias, None
      else:
        input_bias, recurrent_bias = array_ops.unstack(quantized_bias)

    if self.implementation == 1:
      if 0. < self.dropout < 1.:
        inputs_z = inputs * dp_mask[0]
        inputs_r = inputs * dp_mask[1]
        inputs_h = inputs * dp_mask[2]
      else:
        inputs_z = inputs
        inputs_r = inputs
        inputs_h = inputs

      x_z = K.dot(inputs_z, quantized_kernel[:, :self.units])
      x_r = K.dot(inputs_r, quantized_kernel[:, self.units:self.units * 2])
      x_h = K.dot(inputs_h, quantized_kernel[:, self.units * 2:])

      if self.use_bias:
        x_z = K.bias_add(x_z, input_bias[:self.units])
        x_r = K.bias_add(x_r, input_bias[self.units: self.units * 2])
        x_h = K.bias_add(x_h, input_bias[self.units * 2:])

      if 0. < self.recurrent_dropout < 1.:
        h_tm1_z = h_tm1 * rec_dp_mask[0]
        h_tm1_r = h_tm1 * rec_dp_mask[1]
        h_tm1_h = h_tm1 * rec_dp_mask[2]
      else:
        h_tm1_z = h_tm1
        h_tm1_r = h_tm1
        h_tm1_h = h_tm1

      recurrent_z = K.dot(h_tm1_z, quantized_recurrent[:, :self.units])
      recurrent_r = K.dot(h_tm1_r,
                          quantized_recurrent[:, self.units:self.units * 2])
      if self.reset_after and self.use_bias:
        recurrent_z = K.bias_add(recurrent_z, recurrent_bias[:self.units])
        recurrent_r = K.bias_add(recurrent_r,
                                 recurrent_bias[self.units:self.units * 2])

      z = self.recurrent_activation(x_z + recurrent_z)
      r = self.recurrent_activation(x_r + recurrent_r)

      # reset gate applied after/before matrix multiplication
      if self.reset_after:
        recurrent_h = K.dot(h_tm1_h, quantized_recurrent[:, self.units * 2:])
        if self.use_bias:
          recurrent_h = K.bias_add(recurrent_h, recurrent_bias[self.units * 2:])
        recurrent_h = r * recurrent_h
      else:
        recurrent_h = K.dot(r * h_tm1_h,
                            quantized_recurrent[:, self.units * 2:])

      hh = self.activation(x_h + recurrent_h)
    else:
      if 0. < self.dropout < 1.:
        inputs = inputs * dp_mask[0]

      # inputs projected by all gate matrices at once
      matrix_x = K.dot(inputs, quantized_kernel)
      if self.use_bias:
        # biases: bias_z_i, bias_r_i, bias_h_i
        matrix_x = K.bias_add(matrix_x, input_bias)

      x_z, x_r, x_h = array_ops.split(matrix_x, 3, axis=-1)

      if self.reset_after:
        # hidden state projected by all gate matrices at once
        matrix_inner = K.dot(h_tm1, quantized_recurrent)
        if self.use_bias:
          matrix_inner = K.bias_add(matrix_inner, recurrent_bias)
      else:
        # hidden state projected separately for update/reset and new
        matrix_inner = K.dot(h_tm1, quantized_recurrent[:, :2 * self.units])

      recurrent_z, recurrent_r, recurrent_h = array_ops.split(
          matrix_inner, [self.units, self.units, -1], axis=-1)

      z = self.recurrent_activation(x_z + recurrent_z)
      r = self.recurrent_activation(x_r + recurrent_r)

      if self.reset_after:
        recurrent_h = r * recurrent_h
      else:
        recurrent_h = K.dot(r * h_tm1,
                            quantized_recurrent[:, 2 * self.units:])

      hh = self.activation(x_h + recurrent_h)
    # previous and candidate state mixed by update gate
    h = z * h_tm1 + (1 - z) * hh
    return h, [h]

  def get_config(self):
    config = {
        "kernel_quantizer":
            constraints.serialize(self.kernel_quantizer_internal),
        "recurrent_quantizer":
            constraints.serialize(self.recurrent_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal),
        "state_quantizer":
            constraints.serialize(self.state_quantizer_internal)
    }
    base_config = super(QGRUCell, self).get_config()
    return dict(list(base_config.items()) + list(config.items())) class QBidirectional(Bidirectional):
  """
  Class for the QBidirecitonal wrapper.

  Most of these parameters follow the implementation of Bidirectional in
  Keras.

  We refer the reader to the documentation of Bidirectional in Keras for the
  other parameters.

  """
  def get_quantizers(self):
    """
    Returns quantizers in the order they were created.
    """
    return self.forward_layer.get_quantizers() + self.backward_layer.get_quantizers()

  @property
  def activation(self):
    return self.layer.activation

  def get_quantization_config(self):
    return {
      "layer" : self.layer.get_quantization_config(),
      "backward_layer" : self.backward_layer.get_quantization_config()
    } class QSeparableConv1D(SeparableConv1D, PrunableLayer):
  """Depthwise separable 1D convolution."""

  # most of these parameters follow the implementation of SeparableConv1D
  # in Keras, with the exception of depthwise_quantizer, pointwise_quantizer
  # and bias_quantizer.
  #
  # depthwise_quantizer: quantizer function/class for depthwise spatial kernel
  # pointwise_quantizer: quantizer function/class for pointwise kernel
  # bias_quantizer: quantizer function/class for bias
  #
  # we refer the reader to the documentation of SeparableConv1D in Keras for
  # the other parameters.
  #

  def __init__(self,
               filters,
               kernel_size,
               strides=1,
               padding='valid',
               data_format=None,
               dilation_rate=1,
               depth_multiplier=1,
               activation=None,
               use_bias=True,
               depthwise_initializer='glorot_uniform',
               pointwise_initializer='glorot_uniform',
               bias_initializer='zeros',
               depthwise_regularizer=None,
               pointwise_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               depthwise_constraint=None,
               pointwise_constraint=None,
               bias_constraint=None,
               depthwise_quantizer=None,
               pointwise_quantizer=None,
               bias_quantizer=None,
               **kwargs):

    self.depthwise_quantizer = depthwise_quantizer
    self.pointwise_quantizer = pointwise_quantizer
    self.bias_quantizer = bias_quantizer

    self.depthwise_quantizer_internal = get_quantizer(self.depthwise_quantizer)
    self.pointwise_quantizer_internal = get_quantizer(self.pointwise_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)

    # optimize parameter set to "auto" scaling mode if possible
    if hasattr(self.depthwise_quantizer_internal, "_set_trainable_parameter"):
      self.depthwise_quantizer_internal._set_trainable_parameter()

    if hasattr(self.pointwise_quantizer_internal, "_set_trainable_parameter"):
      self.pointwise_quantizer_internal._set_trainable_parameter()

    self.quantizers = [
        self.depthwise_quantizer_internal, self.pointwise_quantizer_internal,
        self.bias_quantizer_internal
    ]

    depthwise_constraint, depthwise_initializer = (
        get_auto_range_constraint_initializer(self.depthwise_quantizer_internal,
                                              depthwise_constraint,
                                              depthwise_initializer))

    pointwise_constraint, pointwise_initializer = (
        get_auto_range_constraint_initializer(self.pointwise_quantizer_internal,
                                              pointwise_constraint,
                                              pointwise_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))

    if activation is not None:
      activation = get_quantizer(activation)

    super(QSeparableConv1D, self).__init__(
        filters=filters,
        kernel_size=kernel_size,
        strides=strides,
        padding=padding,
        data_format=data_format,
        dilation_rate=dilation_rate,
        depth_multiplier=depth_multiplier,
        activation=activation,
        use_bias=use_bias,
        depthwise_initializer=initializers.get(depthwise_initializer),
        pointwise_initializer=initializers.get(pointwise_initializer),
        bias_initializer=initializers.get(bias_initializer),
        depthwise_regularizer=regularizers.get(depthwise_regularizer),
        pointwise_regularizer=regularizers.get(pointwise_regularizer),
        bias_regularizer=regularizers.get(bias_regularizer),
        activity_regularizer=regularizers.get(activity_regularizer),
        depthwise_constraint=constraints.get(depthwise_constraint),
        pointwise_constraint=constraints.get(pointwise_constraint),
        bias_constraint=constraints.get(bias_constraint),
        **kwargs)

  def call(self, inputs):
    if self.padding == 'causal':
      inputs = array_ops.pad(inputs, self._compute_causal_padding())
    
    spatial_start_dim = 1 if self.data_format == 'channels_last' else 2

    # Explicitly broadcast inputs and kernels to 4D.
    inputs = array_ops.expand_dims(inputs, spatial_start_dim)
    depthwise_kernel = array_ops.expand_dims(self.depthwise_kernel, 0)
    pointwise_kernel = array_ops.expand_dims(self.pointwise_kernel, 0)
    dilation_rate = (1,) + self.dilation_rate

    if self.padding == 'causal':
      op_padding = 'valid'
    else:
      op_padding = self.padding

    if self.depthwise_quantizer:
      quantized_depthwise_kernel = self.depthwise_quantizer_internal(
          depthwise_kernel)
    else:
      quantized_depthwise_kernel = depthwise_kernel

    if self.pointwise_quantizer:
      quantized_pointwise_kernel = self.pointwise_quantizer_internal(
          pointwise_kernel)
    else:
      quantized_pointwise_kernel = pointwise_kernel

    outputs = tf.keras.backend.separable_conv2d(
        inputs,
        quantized_depthwise_kernel,
        quantized_pointwise_kernel,
        strides=self.strides * 2,
        padding=op_padding,
        dilation_rate=dilation_rate,
        data_format=self.data_format)

    if self.use_bias:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias

      outputs = tf.keras.backend.bias_add(
          outputs,
          quantized_bias,
          data_format=self.data_format)

    outputs = array_ops.squeeze(outputs, [spatial_start_dim])

    if self.activation is not None:
      return self.activation(outputs)
    return outputs

  def get_config(self):
    config = {
        "depthwise_quantizer":
            constraints.serialize(self.depthwise_quantizer_internal),
        "pointwise_quantizer":
            constraints.serialize(self.pointwise_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal)
    }
    base_config = super(QSeparableConv1D, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantizers(self):
    return self.quantizers

  def get_prunable_weights(self):
    return [self.depthwise_kernel, self.pointwise_kernel] class QSeparableConv2D(SeparableConv2D, PrunableLayer):
  """Depthwise separable 2D convolution."""

  # most of these parameters follow the implementation of SeparableConv2D
  # in Keras, with the exception of depthwise_quantizer, pointwise_quantizer
  # and bias_quantizer.
  #
  # depthwise_quantizer: quantizer function/class for depthwise spatial kernel
  # pointwise_quantizer: quantizer function/class for pointwise kernel
  # bias_quantizer: quantizer function/class for bias
  #
  # we refer the reader to the documentation of SeparableConv2D in Keras for
  # the other parameters.
  #

  def __init__(self,
               filters,
               kernel_size,
               strides=(1, 1),
               padding='valid',
               data_format=None,
               dilation_rate=(1, 1),
               depth_multiplier=1,
               activation=None,
               use_bias=True,
               depthwise_initializer='glorot_uniform',
               pointwise_initializer='glorot_uniform',
               bias_initializer='zeros',
               depthwise_regularizer=None,
               pointwise_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               depthwise_constraint=None,
               pointwise_constraint=None,
               bias_constraint=None,
               depthwise_quantizer=None,
               pointwise_quantizer=None,
               bias_quantizer=None,
               **kwargs):

    self.depthwise_quantizer = depthwise_quantizer
    self.pointwise_quantizer = pointwise_quantizer
    self.bias_quantizer = bias_quantizer

    self.depthwise_quantizer_internal = get_quantizer(self.depthwise_quantizer)
    self.pointwise_quantizer_internal = get_quantizer(self.pointwise_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)

    # optimize parameter set to "auto" scaling mode if possible
    if hasattr(self.depthwise_quantizer_internal, "_set_trainable_parameter"):
      self.depthwise_quantizer_internal._set_trainable_parameter()

    if hasattr(self.pointwise_quantizer_internal, "_set_trainable_parameter"):
      self.pointwise_quantizer_internal._set_trainable_parameter()

    self.quantizers = [
        self.depthwise_quantizer_internal, self.pointwise_quantizer_internal,
        self.bias_quantizer_internal
    ]

    depthwise_constraint, depthwise_initializer = (
        get_auto_range_constraint_initializer(self.depthwise_quantizer_internal,
                                              depthwise_constraint,
                                              depthwise_initializer))

    pointwise_constraint, pointwise_initializer = (
        get_auto_range_constraint_initializer(self.pointwise_quantizer_internal,
                                              pointwise_constraint,
                                              pointwise_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))

    if activation is not None:
      activation = get_quantizer(activation)

    super(QSeparableConv2D, self).__init__(
        filters=filters,
        kernel_size=kernel_size,
        strides=strides,
        padding=padding,
        data_format=data_format,
        dilation_rate=dilation_rate,
        depth_multiplier=depth_multiplier,
        activation=activation,
        use_bias=use_bias,
        depthwise_initializer=initializers.get(depthwise_initializer),
        pointwise_initializer=initializers.get(pointwise_initializer),
        bias_initializer=initializers.get(bias_initializer),
        depthwise_regularizer=regularizers.get(depthwise_regularizer),
        pointwise_regularizer=regularizers.get(pointwise_regularizer),
        bias_regularizer=regularizers.get(bias_regularizer),
        activity_regularizer=regularizers.get(activity_regularizer),
        depthwise_constraint=constraints.get(depthwise_constraint),
        pointwise_constraint=constraints.get(pointwise_constraint),
        bias_constraint=constraints.get(bias_constraint),
        **kwargs)

  def call(self, inputs):
    # Apply the actual ops.
    if self.depthwise_quantizer:
      quantized_depthwise_kernel = self.depthwise_quantizer_internal(
          self.depthwise_kernel)
    else:
      quantized_depthwise_kernel = self.depthwise_kernel

    if self.pointwise_quantizer:
      quantized_pointwise_kernel = self.pointwise_quantizer_internal(
          self.pointwise_kernel)
    else:
      quantized_pointwise_kernel = self.pointwise_kernel

    outputs = tf.keras.backend.separable_conv2d(
        inputs,
        quantized_depthwise_kernel,
        quantized_pointwise_kernel,
        strides=self.strides,
        padding=self.padding,
        dilation_rate=self.dilation_rate,
        data_format=self.data_format)

    if self.use_bias:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias

      outputs = tf.keras.backend.bias_add(
          outputs,
          quantized_bias,
          data_format=self.data_format)

    if self.activation is not None:
      return self.activation(outputs)
    return outputs

  def get_config(self):
    config = {
        "depthwise_quantizer":
            constraints.serialize(self.depthwise_quantizer_internal),
        "pointwise_quantizer":
            constraints.serialize(self.pointwise_quantizer_internal),
        "bias_quantizer":
            constraints.serialize(self.bias_quantizer_internal)
    }
    base_config = super(QSeparableConv2D, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantizers(self):
    return self.quantizers

  def get_prunable_weights(self):
    return [self.depthwise_kernel, self.pointwise_kernel] class QDepthwiseConv2D(DepthwiseConv2D, PrunableLayer):
  """Creates quantized depthwise conv2d. Copied from mobilenet."""

  # most of these parameters follow the implementation of DepthwiseConv2D
  # in Keras, # with the exception of depthwise_range, bias_range,
  # depthwise_quantizer # and bias_quantizer, and kernel_initializer.
  #
  # depthwise_quantizer: quantizer function/class for kernel
  # bias_quantizer: quantizer function/class for bias
  # depthwise_range/bias_ranger: for quantizer functions whose values
  #   can go over [-1,+1], these values are used to set the clipping
  #   value of kernels and biases, respectively, instead of using the
  #   constraints specified by the user.
  #
  # we refer the reader to the documentation of DepthwiseConv2D in Keras for the
  # other parameters.
  #

  def __init__(self,
               kernel_size,
               strides=(1, 1),
               padding="VALID",
               depth_multiplier=1,
               data_format=None,
               activation=None,
               use_bias=True,
               depthwise_initializer="he_normal",
               bias_initializer="zeros",
               depthwise_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               depthwise_constraint=None,
               bias_constraint=None,
               dilation_rate=(1, 1),
               depthwise_quantizer=None,
               bias_quantizer=None,
               depthwise_range=None,
               bias_range=None,
               **kwargs):

    if depthwise_range is not None:
      warnings.warn("depthwise_range is deprecated in QDepthwiseConv2D layer.")

    if bias_range is not None:
      warnings.warn("bias_range is deprecated in QDepthwiseConv2D layer.")

    self.depthwise_range = depthwise_range
    self.bias_range = bias_range

    self.depthwise_quantizer = depthwise_quantizer
    self.bias_quantizer = bias_quantizer

    self.depthwise_quantizer_internal = get_quantizer(self.depthwise_quantizer)
    self.bias_quantizer_internal = get_quantizer(self.bias_quantizer)

    # optimize parameter set to "auto" scaling mode if possible
    if hasattr(self.depthwise_quantizer_internal, "_set_trainable_parameter"):
      self.depthwise_quantizer_internal._set_trainable_parameter()

    self.quantizers = [
        self.depthwise_quantizer_internal, self.bias_quantizer_internal
    ]

    depthwise_constraint, depthwise_initializer = (
        get_auto_range_constraint_initializer(self.depthwise_quantizer_internal,
                                              depthwise_constraint,
                                              depthwise_initializer))

    if use_bias:
      bias_constraint, bias_initializer = (
          get_auto_range_constraint_initializer(self.bias_quantizer_internal,
                                                bias_constraint,
                                                bias_initializer))
    if activation is not None:
      activation = get_quantizer(activation)

    super(QDepthwiseConv2D, self).__init__(
        kernel_size=kernel_size,
        strides=strides,
        padding=padding,
        data_format=data_format,
        activation=activation,
        use_bias=use_bias,
        depthwise_regularizer=depthwise_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        depth_multiplier=depth_multiplier,
        depthwise_initializer=depthwise_initializer,
        bias_initializer=bias_initializer,
        depthwise_constraint=depthwise_constraint,
        bias_constraint=bias_constraint,
        dilation_rate=dilation_rate,
        **kwargs)

  def build(self, input_shape):
    if len(input_shape) < 4:
      raise ValueError(
          "Inputs to `QDepthwiseConv2D` should have rank 4. "
          "Received input shape:", str(input_shape))
    if self.data_format == "channels_first":
      channel_axis = 1
    else:
      channel_axis = 3
    if input_shape[channel_axis] is None:
      raise ValueError("The channel dimension of the inputs to "
                       "`QDepthwiseConv2D` "
                       "should be defined. Found `None`.")
    input_dim = int(input_shape[channel_axis])
    depthwise_kernel_shape = (self.kernel_size[0], self.kernel_size[1],
                              input_dim, self.depth_multiplier)

    self.depthwise_kernel = self.add_weight(
        shape=depthwise_kernel_shape,
        initializer=self.depthwise_initializer,
        name="depthwise_kernel",
        regularizer=self.depthwise_regularizer,
        constraint=self.depthwise_constraint)

    if self.use_bias:
      self.bias = self.add_weight(
          shape=(input_dim * self.depth_multiplier,),
          initializer=self.bias_initializer,
          name="bias",
          regularizer=self.bias_regularizer,
          constraint=self.bias_constraint)
    else:
      self.bias = None
    # Set input spec.
    self.input_spec = InputSpec(ndim=4, axes={channel_axis: input_dim})
    self.built = True

  def call(self, inputs, training=None):
    if self.depthwise_quantizer:
      quantized_depthwise_kernel = (
          self.depthwise_quantizer_internal(self.depthwise_kernel))
    else:
      quantized_depthwise_kernel = self.depthwise_kernel
    outputs = tf.keras.backend.depthwise_conv2d(
        inputs,
        quantized_depthwise_kernel,
        strides=self.strides,
        padding=self.padding,
        dilation_rate=self.dilation_rate,
        data_format=self.data_format)

    if self.use_bias:
      if self.bias_quantizer:
        quantized_bias = self.bias_quantizer_internal(self.bias)
      else:
        quantized_bias = self.bias
      outputs = tf.keras.backend.bias_add(
          outputs, quantized_bias, data_format=self.data_format)

    if self.activation is not None:
      return self.activation(outputs)

    return outputs

  def get_config(self):
    config = super(QDepthwiseConv2D, self).get_config()
    config.pop("filters", None)
    config.pop("kernel_initializer", None)
    config.pop("kernel_regularizer", None)
    config.pop("kernel_constraint", None)
    config["depth_multiplier"] = self.depth_multiplier
    config["depthwise_initializer"] = initializers.serialize(
        self.depthwise_initializer)
    config["depthwise_regularizer"] = regularizers.serialize(
        self.depthwise_regularizer)
    config["depthwise_constraint"] = constraints.serialize(
        self.depthwise_constraint)
    config["depthwise_quantizer"] = constraints.serialize(
        self.depthwise_quantizer_internal)
    config["bias_quantizer"] = constraints.serialize(
        self.bias_quantizer_internal)
    config["depthwise_range"] = self.depthwise_range
    config["bias_range"] = self.bias_range
    return config

  def get_quantization_config(self):
    return {
        "depthwise_quantizer_internal":
            str(self.depthwise_quantizer_internal),
        "bias_quantizer":
            str(self.bias_quantizer_internal),
        "activation":
            str(self.activation),
        "filters" : str(self.filters)
    }

  def get_quantizers(self):
    return self.quantizers

  def get_prunable_weights(self):
    return [self.depthwise_kernel] class QBatchNormalization(BatchNormalization, PrunableLayer):
  """Quantized Batch Normalization layer.
  For training, mean and variance are not quantized.
  For inference, the quantized moving mean and moving variance are used.

  output = (x - mean) / sqrt(var + epsilon) * quantized_gamma + quantized_beta

  """

  def __init__(
      self,
      axis=-1,
      momentum=0.99,
      epsilon=1e-3,
      center=True,
      scale=True,
      activation=None,
      beta_initializer='zeros',
      gamma_initializer='ones',
      moving_mean_initializer='zeros',
      moving_variance_initializer='ones',
      beta_regularizer=None,
      gamma_regularizer=None,
      beta_quantizer='quantized_po2(5)',
      gamma_quantizer='quantized_relu_po2(6, 2048)',
      mean_quantizer='quantized_po2(5)',
      variance_quantizer='quantized_relu_po2(6, quadratic_approximation=True)',
      inverse_quantizer=None,
      gamma_constraint=None,
      beta_constraint=None,
      # use quantized_po2 and enforce quadratic approximation
      # to get an even exponent for sqrt
      beta_range=None,
      gamma_range=None,
      **kwargs):

    if gamma_range is not None:
      warnings.warn('gamma_range is deprecated in QBatchNormalization layer.')

    if beta_range is not None:
      warnings.warn('beta_range is deprecated in QBatchNormalization layer.')

    self.gamma_range = gamma_range
    self.beta_range = beta_range
    self.activation = activation

    self.beta_quantizer = beta_quantizer
    self.gamma_quantizer = gamma_quantizer
    self.mean_quantizer = mean_quantizer
    self.variance_quantizer = variance_quantizer
    self.inverse_quantizer = inverse_quantizer

    if self.inverse_quantizer is not None:
      assert self.variance_quantizer is None and self.gamma_quantizer is None, (
          'If using the inverse quantizer, the gamma and variance quantizers '
          'should not be used in order to avoid quantizing a value twice.')

    self.beta_quantizer_internal = get_quantizer(self.beta_quantizer)
    self.gamma_quantizer_internal = get_quantizer(self.gamma_quantizer)
    self.mean_quantizer_internal = get_quantizer(self.mean_quantizer)
    self.variance_quantizer_internal = get_quantizer(self.variance_quantizer)
    self.inverse_quantizer_internal = get_quantizer(self.inverse_quantizer)

    if hasattr(self.gamma_quantizer_internal, '_set_trainable_parameter'):
      self.gamma_quantizer_internal._set_trainable_parameter()
    if hasattr(self.variance_quantizer_internal, '_set_trainable_parameter'):
      self.variance_quantizer_internal._set_trainable_parameter()

    self.quantizers = [
        self.gamma_quantizer_internal,
        self.beta_quantizer_internal,
        self.mean_quantizer_internal,
        self.variance_quantizer_internal,
        self.inverse_quantizer_internal
    ]

    if scale and self.gamma_quantizer:
      gamma_constraint, gamma_initializer = (
          get_auto_range_constraint_initializer(
              self.gamma_quantizer_internal,
              gamma_constraint,
              gamma_initializer)
      )

    if center and self.beta_quantizer:
      beta_constraint, beta_initializer = (
          get_auto_range_constraint_initializer(
              self.beta_quantizer_internal,
              beta_constraint,
              beta_initializer)
      )

    if kwargs.get('fused', None):
      warnings.warn('batch normalization fused is disabled '
                    'in qkeras qnormalization.py.')
      del kwargs['fused']

    if kwargs.get('renorm', None):
      warnings.warn('batch normalization renorm is disabled '
                    'in qkeras qnormalization.py.')
      del kwargs['renorm']

    if kwargs.get('virtual_batch_size', None):
      warnings.warn('batch normalization virtual_batch_size is disabled '
                    'in qkeras qnormalization.py.')
      del kwargs['virtual_batch_size']

    if kwargs.get('adjustment', None):
      warnings.warn('batch normalization adjustment is disabled '
                    'in qkeras qnormalization.py.')
      del kwargs['adjustment']

    super(QBatchNormalization, self).__init__(
        axis=axis,
        momentum=momentum,
        epsilon=epsilon,
        center=center,
        scale=scale,
        beta_initializer=beta_initializer,
        gamma_initializer=gamma_initializer,
        moving_mean_initializer=moving_mean_initializer,
        moving_variance_initializer=moving_variance_initializer,
        beta_regularizer=beta_regularizer,
        gamma_regularizer=gamma_regularizer,
        beta_constraint=beta_constraint,
        gamma_constraint=gamma_constraint,
        fused=False,
        renorm=False,
        virtual_batch_size=None,
        adjustment=None,
        **kwargs)

  def call(self, inputs, training=None):
    if self.scale and self.gamma_quantizer:
      quantized_gamma = self.gamma_quantizer_internal(self.gamma)
    else:
      quantized_gamma = self.gamma

    if self.center and self.beta_quantizer:
      quantized_beta = self.beta_quantizer_internal(self.beta)
    else:
      quantized_beta = self.beta

    if self.mean_quantizer:
      quantized_moving_mean = self.mean_quantizer_internal(self.moving_mean)
    else:
      quantized_moving_mean = self.moving_mean

    if self.variance_quantizer:
      quantized_moving_variance = self.variance_quantizer_internal(
          self.moving_variance)
    else:
      quantized_moving_variance = self.moving_variance

    training = self._get_training_value(training)

    # Compute the axes along which to reduce the mean / variance
    input_shape = inputs.shape
    ndims = len(input_shape)
    reduction_axes = [i for i in range(ndims) if i not in self.axis]

    # Broadcasting only necessary for single-axis batch norm where the axis is
    # not the last dimension
    broadcast_shape = [1] * ndims
    broadcast_shape[self.axis[0]] = input_shape.dims[self.axis[0]].value
    def _broadcast(v):
      if (v is not None and len(v.shape) != ndims and
          reduction_axes != list(range(ndims - 1))):
        return array_ops.reshape(v, broadcast_shape)
      return v

    scale, offset = _broadcast(quantized_gamma), _broadcast(quantized_beta)

    # Determine a boolean value for `training`: could be True, False, or None.
    training_value = tf_utils.smart_constant_value(training)
    if training_value == False:  # pylint: disable=singleton-comparison,g-explicit-bool-comparison
      quantized_mean, quantized_variance = (quantized_moving_mean,
                                            quantized_moving_variance)
    else:
      # Some of the computations here are not necessary when training==False
      # but not a constant. However, this makes the code simpler.
      keep_dims = len(self.axis) > 1
      mean, variance = self._moments(
          math_ops.cast(inputs, self._param_dtype),
          reduction_axes,
          keep_dims=keep_dims)

      moving_mean = self.moving_mean
      moving_variance = self.moving_variance

      mean = tf_utils.smart_cond(
          training, lambda: mean, lambda: ops.convert_to_tensor(moving_mean))
      variance = tf_utils.smart_cond(
          training,
          lambda: variance,
          lambda: ops.convert_to_tensor(moving_variance))

      new_mean, new_variance = mean, variance

      if self.mean_quantizer:
        quantized_mean = self.mean_quantizer_internal(mean)
      else:
        quantized_mean = mean

      if self.variance_quantizer:
        quantized_variance = self.variance_quantizer_internal(variance)
      else:
        quantized_variance = variance

      if self._support_zero_size_input():
        inputs_size = array_ops.size(inputs)
      else:
        inputs_size = None

      def _do_update(var, value):
        """Compute the updates for mean and variance."""
        return self._assign_moving_average(var, value, self.momentum,
                                           inputs_size)

      def mean_update():
        true_branch = lambda: _do_update(self.moving_mean, new_mean)
        false_branch = lambda: self.moving_mean
        return tf_utils.smart_cond(training, true_branch, false_branch)

      def variance_update():
        """Update the moving variance."""
        true_branch = lambda: _do_update(self.moving_variance, new_variance)
        false_branch = lambda: self.moving_variance
        return tf_utils.smart_cond(training, true_branch, false_branch)

      self.add_update(mean_update)
      self.add_update(variance_update)

    quantized_mean = _broadcast(math_ops.cast(quantized_mean, inputs.dtype))
    quantized_variance = _broadcast(
        math_ops.cast(quantized_variance, inputs.dtype))
    if offset is not None:
      offset = math_ops.cast(offset, inputs.dtype)
    if scale is not None:
      scale = math_ops.cast(scale, inputs.dtype)

    # Calculate and quantize the inverse
    inv = math_ops.rsqrt(quantized_variance + self.epsilon)
    if scale is not None:
      inv *= scale
    if self.inverse_quantizer_internal is not None:
      inv = self.inverse_quantizer_internal(inv)

    # Calculate the forward pass of the BN
    outputs = inputs * math_ops.cast(inv, inputs.dtype) + math_ops.cast(
        offset - quantized_mean * inv
        if offset is not None else -quantized_mean * inv, inputs.dtype)

    # If some components of the shape got lost due to adjustments, fix that.
    outputs.set_shape(input_shape)

    return outputs

  def get_config(self):
    config = {
        'axis': self.axis,
        'momentum': self.momentum,
        'epsilon': self.epsilon,
        'center': self.center,
        'scale': self.scale,
        'beta_quantizer':
            constraints.serialize(self.beta_quantizer_internal),
        'gamma_quantizer':
            constraints.serialize(self.gamma_quantizer_internal),
        'mean_quantizer':
            constraints.serialize(self.mean_quantizer_internal),
        'variance_quantizer':
            constraints.serialize(self.variance_quantizer_internal),
        'beta_initializer': initializers.serialize(self.beta_initializer),
        'gamma_initializer': initializers.serialize(self.gamma_initializer),
        'moving_mean_initializer':
            initializers.serialize(self.moving_mean_initializer),
        'moving_variance_initializer':
            initializers.serialize(self.moving_variance_initializer),
        'inverse_quantizer':
            initializers.serialize(self.inverse_quantizer_internal),
        'beta_regularizer': regularizers.serialize(self.beta_regularizer),
        'gamma_regularizer': regularizers.serialize(self.gamma_regularizer),
        'beta_constraint': constraints.serialize(self.beta_constraint),
        'gamma_constraint': constraints.serialize(self.gamma_constraint),
        'beta_range': self.beta_range,
        'gamma_range': self.gamma_range,
    }
    base_config = super(QBatchNormalization, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def compute_output_shape(self, input_shape):
    return input_shape

  def get_quantizers(self):
    return self.quantizers

  def get_prunable_weights(self):
    return [] class QGlobalAveragePooling2D(GlobalAveragePooling2D):
  """Computes the quantized version of GlobalAveragePooling2D."""

  def __init__(self, data_format=None,
               average_quantizer=None,
               activation=None,
               **kwargs):

    self.average_quantizer = average_quantizer
    self.average_quantizer_internal = get_quantizer(self.average_quantizer)
    self.quantizers = [self.average_quantizer_internal]

    if activation is not None:
      self.activation = get_quantizer(activation)
    else:
      self.activation = activation

    super(QGlobalAveragePooling2D, self).__init__(
        data_format=data_format, **kwargs)

  def compute_pooling_area(self, input_shape):
    input_shape = input_shape.as_list()
    if self.data_format == "channels_last":
      return input_shape[1] * input_shape[2]
    else:
      return input_shape[2] * input_shape[3]

  def call(self, inputs):
    """Performs quantized GlobalAveragePooling followed by QActivation.

    Since there is no specific parameter for averaging op, we couldn't apply
    averaging quantizer to the averaging op. We have two options:
    1. we perform our own average as sum first then multiply with the
       inversion
       of the division factor: sum(x) * quantize(1/pool_area)
    2. first, we call keras version of averaging first:
       y1 = keras_global_average(x)
       then multiply it with the denominator(pool_area) used by averaging:
       y2 = y1 * pool_area
       Last, y3 = y2 * quantize(1/ pool_area)
    Our numerical anaysis suggests negligible error between 1 and 2. Therefore
    we use option #2 here for the simplicity of implementation.

    """

    x = super(QGlobalAveragePooling2D, self).call(inputs)

    if self.average_quantizer:
      pool_area = self.compute_pooling_area(input_shape=inputs.shape)
      # Reverts the division results
      x = x * pool_area
      # Quantizes the multiplication factor
      mult_factor = 1.0 / pool_area
      q_mult_factor = self.average_quantizer_internal(mult_factor)
      x = x * q_mult_factor

    if self.activation is not None:
      return self.activation(x)
    return x

  def get_config(self):
    config = {
        "average_quantizer":
            constraints.serialize(self.average_quantizer_internal),
        "activation": constraints.serialize(self.activation),
    }
    base_config = super(QGlobalAveragePooling2D, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def get_quantization_config(self):
    return {
        "average_quantizer":
            str(self.average_quantizer_internal),
        "activation":
            str(self.activation)
    }

  def get_quantizers(self):
    return self.quantizers SOURCE = -1
SINK = -2
class WrongInputQuantizerError(ValueError):
def GraphRemoveNode(graph, v):
def GraphRemoveNodeWithNodeType(graph, node_type):
def  GraphAddHiddenInputLayer(model, graph, input_quantizer_map):
def GraphAddSingleSourceSingleSink(graph):
def GenerateInputQuantizerList(input_quantizers,
                               inputs_length,
                               default_source_quantizer):
def AddToNodeDict(layer_items,
                  layer,
                  nodes_dict):
def GenerateGraphFromModel(model,
                           input_quantizers,
                           default_source_quantizer):
def GraphGetInputs(graph):
def GraphGetOutputs(graph):
def GraphPropagateActivationsToEdges(graph, debug=False):
def PrintGraph(graph, msg=""):
def CreateGraph(model, input_quantizers=None,
                default_source_quantizer=cfg.default_source_quantizer,
                debug=False):
def GraphUpdateEdge(graph, node_id, quantizer_on_edge): class binary(BaseQuantizer):  # pylint: disable=invalid-name
  """Computes the sign(x) returning a value between -alpha and alpha.

  Although we cannot guarantee E[dL/dy] = E[dL/dx] if we do not use the
  stochastic sampling, we still use the ST approximation.

  Modified from original binary to match QNN implementation.

  Attributes:
    x: tensor to perform sign_through.
    bits: number of bits to perform quantization.
    use_01: if True, return {0,1} instead of {-1,+1}.
    alpha: binary is -alpha or +alpha, or "auto", "auto_po2" to compute
      automatically.
    use_stochastic_rounding: if true, we perform stochastic rounding.

  Returns:
    Computation of sign operation with straight through gradient.
  """

  def __init__(self, use_01=False, alpha=None, use_stochastic_rounding=False):
    super(binary, self).__init__()
    self.use_01 = use_01
    self.bits = 1
    self.alpha = alpha
    self.use_stochastic_rounding = use_stochastic_rounding
    self.default_alpha = 1.0
    self.scale = None

  def __str__(self):
    flags = []
    if self.use_01:
      flags.append("use_01=" + str(int(self.use_01)))
    if self.alpha is not None:
      alpha = str(self.alpha)
      if isinstance(self.alpha, six.string_types):
        alpha = "'" + alpha + "'"
      flags.append("alpha=" + alpha)
    if self.use_stochastic_rounding:
      flags.append(
          "use_stochastic_rounding=" + str(self.use_stochastic_rounding))
    return "binary(" + ",".join(flags) + ")"

  def __call__(self, x):
    if isinstance(self.alpha, six.string_types):
      assert self.alpha in ["auto", "auto_po2"]
    if self.alpha is None:
      scale = self.default_alpha
    elif isinstance(self.alpha, six.string_types):
      scale = 1.0
    elif isinstance(self.alpha, np.ndarray):
      scale = self.alpha
    else:
      scale = float(self.alpha)

    if self.use_stochastic_rounding:
      try:
        len_axis = len(x.shape.as_list())
      except AttributeError:
        len_axis = len(list(x.shape))
      if len_axis == 1:
        axis = None
      elif K.image_data_format() == "channels_last":
        axis = list(range(len_axis - 1))
      else:
        axis = list(range(1, len_axis))

      # if stochastic_round is through, we need to scale
      # number so that the precision is small enough.
      # This is especially important if range of x is very
      # small, which occurs during initialization of weights.
      m = K.max(tf.abs(x), axis=axis, keepdims=True)
      m = tf.where(m > 1.0, tf.ones_like(m), m)
      f = 2 * m

      x = tf_utils.smart_cond(
          K.learning_phase(),
          lambda: f * _round_through(
              x / f, use_stochastic_rounding=True, precision=0.125),
          lambda: x)

    k_sign = tf.sign(x)
    if self.use_stochastic_rounding:
      # in inference, we use a biased "1" for stochastic rounding right now
      k_sign += (1.0 - tf.abs(k_sign)) * tf_utils.smart_cond(
          K.learning_phase(),
          lambda: 2.0 * tf.round(tf.random.uniform(tf.shape(x))) - 1.0,
          lambda: tf.ones_like(tf.shape(x), dtype=K.floatx()))
      # if something still remains, just make it positive for now.
    k_sign += (1.0 - tf.abs(k_sign))
    if self.use_01:
      k_sign = (k_sign + 1.0) / 2.0

    # approximate binary by tanh(x) as it has limited range between -1 and +1.
    if self.alpha is None:
      x = K.tanh(x)

    scale = _get_scale(self.alpha, x, k_sign)
    self.scale = scale
    return x + tf.stop_gradient(-x + scale * k_sign)

  def _set_trainable_parameter(self):
    if self.alpha is None:
      self.alpha = "auto_po2"

  def max(self):
    """Get maximum value that binary class can respresent."""
    if self.alpha is None or isinstance(self.alpha, six.string_types):
      return 1.0
    else:
      return max(1.0, self.alpha)

  def min(self):
    """Get minimum value that binary class can respresent."""
    if self.use_01:
      return 0.0
    elif self.alpha is None or isinstance(self.alpha, six.string_types):
      return -1.0
    else:
      return -max(1.0, self.alpha)

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "use_01": self.use_01,
        "alpha": self.alpha,
        "use_stochastic_rounding": self.use_stochastic_rounding
    }
    return config class bernoulli(BaseQuantizer):  # pylint: disable=invalid-name
  """Computes a Bernoulli sample with probability sigmoid(x).

  This computation uses ST approximation.

  To do that, we compute sigmoid(x) and a random sample z ~ U[0,1]. As
  p in [0,1] and z in [0,1], p - z in [-1,1]. However, -1 will
  never appear because to get -1 we would need sigmoid(-inf) - z == 1.
  As a result, the range will be in practical terms [0,1].

  The noise introduced by z can be seen as a regularizer to the weights W of
  y = Wx as y = Wx + Wz for some noise z with mean mu(z) and var(z). As a
  result, W**2 var(z) to the variance of y, which has the same effect as a
  regularizer on L2 with lambda = var(z), as presented in Hinton"s Coursera
  Lecture 9c.

  Remember that E[dL/dy] = E[dL/dx] once we add stochastic sampling.

  Attributes:
    alpha: allows one to specify multiplicative factor for number generation
      of "auto" or "auto_po2".
    temperature: amplifier factor for sigmoid function, making stochastic
      less stochastic as it moves away from 0.
    use_real_sigmoid: use real sigmoid for probability.

  Returns:
    Computation of round with stochastic sampling with straight through
    gradient.
  """

  def __init__(self, alpha=None, temperature=6.0, use_real_sigmoid=True):
    super(bernoulli, self).__init__()
    self.alpha = alpha
    self.bits = 1
    self.temperature = temperature
    self.use_real_sigmoid = use_real_sigmoid
    self.default_alpha = 1.0
    self.scale = None

  def __str__(self):
    flags = []
    if self.alpha is not None:
      alpha = str(self.alpha)
      if isinstance(self.alpha, six.string_types):
        alpha = "'" + alpha + "'"
      flags.append("alpha=" + alpha)
    if self.temperature != 6.0:
      flags.append("temperature=" + str(self.temperature))
    if not self.use_real_sigmoid:
      flags.append("use_real_sigmoid=" + str(int(self.use_real_sigmoid)))
    return "bernoulli(" + ",".join(flags) + ")"

  def __call__(self, x):
    if isinstance(self.alpha, six.string_types):
      assert self.alpha in ["auto", "auto_po2"]

    if isinstance(self.alpha, six.string_types):
      len_axis = len(x.shape)

      if len_axis > 1:
        if K.image_data_format() == "channels_last":
          axis = list(range(len_axis - 1))
        else:
          axis = list(range(1, len_axis))
      else:
        axis = [0]

      std = K.std(x, axis=axis, keepdims=True) + K.epsilon()
    else:
      std = 1.0

    if self.use_real_sigmoid:
      p = tf.keras.backend.sigmoid(self.temperature * x / std)
    else:
      p = _sigmoid(self.temperature * x/std)
    r = tf.random.uniform(tf.shape(x))
    q = tf.sign(p - r)
    q += (1.0 - tf.abs(q))
    q = (q + 1.0) / 2.0

    q_non_stochastic = tf.sign(x)
    q_non_stochastic += (1.0 - tf.abs(q_non_stochastic))
    q_non_stochastic = (q_non_stochastic + 1.0) / 2.0

    # if we use non stochastic binary to compute alpha,
    # this function seems to behave better
    scale = _get_scale(self.alpha, x, q_non_stochastic)
    self.scale = scale
    return x + tf.stop_gradient(-x + scale * q)

  def _set_trainable_parameter(self):
    if self.alpha is None:
      self.alpha = "auto_po2"

  def max(self):
    """Get the maximum value bernoulli class can represent."""
    if self.alpha is None or isinstance(self.alpha, six.string_types):
      return 1.0
    else:
      return max(1.0, self.alpha)

  def min(self):
    """Get the minimum value bernoulli class can represent."""
    return 0.0

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {"alpha": self.alpha}
    return config def get_weight_scale(quantizer, x=None):
  """Gets the scales of weights for (stochastic_)binary and ternary quantizers.

  Arguments:
    quantizer: A binary or teneray quantizer class.
    x: A weight tensor.  We keep it here for now for backward compatibility.

  Returns:
    Weight scale per channel for binary and ternary
    quantizers with auto or auto_po2 alpha/threshold.
  """
  if hasattr(quantizer, "scale") and quantizer.scale is not None:
    return K.eval(quantizer.scale)
  return 1.0 class quantized_bits(BaseQuantizer):  # pylint: disable=invalid-name
  """Quantizes the number to a number of bits.

  In general, we want to use a quantization function like:

  a = (pow(2,bits) - 1 - 0) / (max(x) - min(x))
  b = -min(x) * a

  in the equation:

  xq = a x + b

  This requires multiplication, which is undesirable. So, we
  enforce weights to be between -1 and 1 (max(x) = 1 and min(x) = -1),
  and separating the sign from the rest of the number as we make this function
  symmetric, thus resulting in the following approximation.

  1) max(x) = +1, min(x) = -1
  2) max(x) = -min(x)

  a = pow(2,bits-1)
  b = 0

  Finally, just remember that to represent the number with sign, the
  largest representation is -pow(2,bits) to pow(2, bits-1)

  Symmetric and keep_negative allow us to generate numbers that are symmetric
  (same number of negative and positive representations), and numbers that
  are positive.

  Note:
    the behavior of quantized_bits is different than Catapult HLS ac_fixed
    or Vivado HLS ap_fixed. For ac_fixed<word_length, integer_lenth, signed>,
    when signed = true, it is equavlent to
    quantized_bits(word_length, integer_length-1, keep_negative=True)

  Attributes:
    bits: number of bits to perform quantization.
    integer: number of bits to the left of the decimal point.
    symmetric: if true, we will have the same number of values for positive
      and negative numbers.
    alpha: a tensor or None, the scaling factor per channel.
      If None, the scaling factor is 1 for all channels.
    keep_negative: if true, we do not clip negative numbers.
    use_stochastic_rounding: if true, we perform stochastic rounding.
    scale_axis: which axis to calculate scale from
    qnoise_factor: float. a scalar from 0 to 1 that represents the level of
      quantization noise to add. This controls the amount of the quantization
      noise to add to the outputs by changing the weighted sum of
      (1 - qnoise_factor)*unquantized_x + qnoise_factor*quantized_x.
    var_name: String or None. A variable name shared between the tf.Variables
      created in the build function. If None, it is generated automatically.
    use_ste: Bool. Whether to use "straight-through estimator" (STE) method or
        not.
    use_variables: Bool. Whether to make the quantizer variables to be dynamic
      tf.Variables or not.

  Returns:
    Function that computes fixed-point quantization with bits.
  """

  def __init__(self,
               bits=8,
               integer=0,
               symmetric=0,
               keep_negative=True,
               alpha=None,
               use_stochastic_rounding=False,
               scale_axis=None,
               qnoise_factor=1.0,
               var_name=None,
               use_ste=True,
               use_variables=False):
    super(quantized_bits, self).__init__()
    self.bits = bits
    self.integer = integer
    self.symmetric = symmetric
    self.keep_negative = keep_negative
    self.alpha = alpha
    self.use_stochastic_rounding = use_stochastic_rounding
    # "auto*" |-> symmetric
    if isinstance(self.alpha, six.string_types):
      self.symmetric = True
    self.scale = None
    self.scale_axis = scale_axis
    self.qnoise_factor = qnoise_factor
    self.use_ste = use_ste
    self.var_name = var_name
    self.use_variables = use_variables

  def __str__(self):
    # Convert Tensors to printable strings by converting to a numpy array and
    # then using regex to remove brackets when there is only one integer bit
    integer_bits = re.sub(
        r"\[(\d)\]", r"\g<1>",
        str(self.integer.numpy() if isinstance(self.integer, tf.Variable
                                              ) else self.integer))

    flags = [str(self.bits), integer_bits, str(int(self.symmetric))]
    if not self.keep_negative:
      flags.append("keep_negative=False")
    if self.alpha:
      alpha = str(self.alpha)
      if isinstance(self.alpha, six.string_types):
        alpha = "'" + alpha + "'"
      flags.append("alpha=" + alpha)
    if self.use_stochastic_rounding:
      flags.append("use_stochastic_rounding=" +
                   str(int(self.use_stochastic_rounding)))
    return "quantized_bits(" + ",".join(flags) + ")"

  def __call__(self, x):
    """Computes fixedpoint quantization of x."""
    if not self.built:
      self.build(var_name=self.var_name, use_variables=self.use_variables)

    x = K.cast_to_floatx(x)

    # quantized_bits with "1" bit becomes a binary implementation.
    unsigned_bits = self.bits - self.keep_negative
    m = K.cast_to_floatx(pow(2, unsigned_bits))
    m_i = K.cast_to_floatx(K.pow(2, self.integer))

    if self.alpha is None:
      scale = 1.0
    elif isinstance(self.alpha, six.string_types):
      # We only deal with the symmetric case right now.
      assert self.symmetric, "Only symmetric quantizers are implemented"
      len_axis = len(x.shape)
      if len_axis > 1:
        axis = _get_scaling_axis(self.scale_axis, len_axis)
      else:
        axis = [0]

      x = x / m_i

      # Using 2's complement, we can represent 2**(bits-1)-1 positive values
      # If we wish to maintain symmetry, we can double 2**(bits-1)-1 to get
      # the total number of possible values we can represent.
      # If symmetry is not enforced, then we can represent (2**bits)-1 values
      # using 2's complement.
      levels = (2**(self.bits-1)-1) * 2 if self.symmetric else (2**self.bits)-1

      scale = (K.max(abs(x), axis=axis, keepdims=True) * 2) / levels

      # If alpha is "auto_po2", then get the "best" po2 scale
      if "po2" in self.alpha:
        scale = K.pow(2.0,
                      tf.math.round(K.log(scale + K.epsilon()) / np.log(2.0)))
        for _ in range(5):
          v = tf.floor(tf.abs(x) / scale + 0.5)
          mask = v < levels / 2
          z = tf.sign(x) * tf.where(mask, v, tf.ones_like(v) * levels / 2)
          scale = _get_scale(alpha="auto_po2", x=x, q=z,
                             scale_axis=self.scale_axis)

      # If alpha is "auto", then get the "best" floating point scale
      elif self.alpha == "auto":
        v = tf.floor(tf.abs(x) / scale + 0.5)
        mask = v < levels / 2
        z = tf.sign(x) * tf.where(mask, v, tf.ones_like(v) * levels / 2)
      else:
        raise ValueError(f"Invalid alpha '{self.alpha}'")

      # z is an integer number, so we must make the scale * m and z / m
      scale = scale * m

      # we will not use "z" right now because of stochastic_rounding
      # this is still under test.

      # if "new" in self.alpha:
      #  z = z / m
      #  self.scale = scale
      #  return x + tf.stop_gradient(-x + scale * z)
      x = m_i * x
      xq = m_i * z / m
      self.scale = scale
      xq = scale * xq

      if self.use_ste:
        return x + tf.stop_gradient(self.qnoise_factor * (-x + xq))
      else:
        return (1 - self.qnoise_factor) * x + tf.stop_gradient(
            self.qnoise_factor * xq)

    else:
      scale = self.alpha

    # quantized_bits with "1" bit becomes a binary implementation.
    if unsigned_bits > 0:
      p = x * m / m_i
      xq = m_i * tf.keras.backend.clip(
          _round_through(p, self.use_stochastic_rounding, precision=1.0),
          self.keep_negative  * (-m + self.symmetric), m - 1) / m
    else:
      xq = tf.sign(x)
      xq += (1.0 - tf.abs(xq))
      if not self.keep_negative:
        xq = (xq + 1.0) / 2.0

    self.scale = scale
    xq = scale * xq

    if self.use_ste:
      return x + tf.stop_gradient(self.qnoise_factor * (-x + xq))
    else:
      return (1 - self.qnoise_factor) * x + tf.stop_gradient(
          self.qnoise_factor * xq)

  def _set_trainable_parameter(self):
    if self.alpha is None:
      self.alpha = "auto_po2"
      self.symmetric = True

  def max(self):
    """Get maximum value that quantized_bits class can represent."""
    unsigned_bits = self.bits - self.keep_negative
    if unsigned_bits > 0:
      return max(
          1.0,
          np.array(
              K.pow(2., K.cast(self.integer, dtype="float32")),
              dtype="float32"))
    else:
      return 1.0

  def min(self):
    """Get minimum value that quantized_bits class can represent."""
    if not self.keep_negative:
      return 0.0
    unsigned_bits = self.bits - self.keep_negative
    if unsigned_bits > 0:
      return -max(
          1.0,
          np.array(
              K.pow(2, K.cast(self.integer, dtype="float32")), dtype="float32"))
    else:
      return -1.0

  def range(self):
    """Returns a list of all values that quantized_bits can represent
    ordered by their binary representation ascending."""
    assert self.symmetric == 0
    assert self.keep_negative
    assert self.alpha is None or self.alpha == 1.0

    x = np.asarray(range(2**self.bits), dtype=np.float32)
    p_and_n = np.where(x >= 2**(self.bits - 1),
                       (x - 2**(self.bits - 1)) - 2**(self.bits - 1), x)
    return p_and_n * np.array(
        K.pow(2.0, -self.bits + K.cast(self.integer, dtype="float32") + 1),
        dtype="float32")

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "bits":
            self.bits,
        "integer":
            self.integer.numpy()
            if isinstance(self.integer, tf.Variable) else self.integer,
        "symmetric":
            self.symmetric,
        "alpha":
            self.alpha,
        "keep_negative":
            self.keep_negative,
        "use_stochastic_rounding":
            self.use_stochastic_rounding,
        "qnoise_factor":
            self.qnoise_factor.numpy() if isinstance(
                self.qnoise_factor, tf.Variable) else self.qnoise_factor
    }
    return config class quantized_relu(BaseQuantizer):  # pylint: disable=invalid-name
  """Computes a quantized relu to a number of bits.

  Modified from:

  [https://github.com/BertMoons/QuantizedNeuralNetworks-Keras-Tensorflow]

  Assume h(x) = +1 with p = sigmoid(x), -1 otherwise, the expected value of
  h(x) is:

  E[h(x)] = +1 P(p <= sigmoid(x)) - 1 P(p > sigmoid(x))
          = +1 P(p <= sigmoid(x)) - 1 ( 1 - P(p <= sigmoid(x)) )
          = 2 P(p <= sigmoid(x)) - 1
          = 2 sigmoid(x) - 1, if p is sampled from a uniform distribution U[0,1]

  If use_sigmoid is 0, we just keep the positive numbers up to
  2**integer * (1 - 2**(-bits)) instead of normalizing them, which is easier
  to implement in hardware.

  Attributes:
    bits: number of bits to perform quantization.
    integer: number of bits to the left of the decimal point.
    use_sigmoid: if true, we apply sigmoid to input to normalize it.
    negative_slope: slope when activation < 0, needs to be power of 2.
    use_stochastic_rounding: if true, we perform stochastic rounding.
    relu_upper_bound: A float representing an upper bound of the unquantized
      relu. If None, we apply relu without the upper bound when
      "is_quantized_clip" is set to false (true by default).
      Note: The quantized relu uses the quantization parameters (bits and
      integer) to upper bound. So it is important to set relu_upper_bound
      appropriately to the quantization parameters. "is_quantized_clip"
      has precedence over "relu_upper_bound" for backward compatibility.
    is_quantized_clip: A boolean representing whether the inputs are clipped to
      the maximum value represented by the quantization parameters. This
      parameter is deprecated, and the default is set to True for backwards
      compatibility. Users are encouraged to use "relu_upper_bound" instead.
    qnoise_factor: float. a scalar from 0 to 1 that represents the level of
      quantization noise to add. This controls the amount of the quantization
      noise to add to the outputs by changing the weighted sum of
      (1 - qnoise_factor)*unquantized_x + qnoise_factor*quantized_x.
    var_name: String or None. A variable name shared between the tf.Variables
      created in the build function. If None, it is generated automatically.
    use_ste: Bool. Whether to use "straight-through estimator" (STE) method or
        not.
    use_variables: Bool. Whether to make the quantizer variables to be dynamic
      tf.Variables or not.

  Returns:
    Function that performs relu + quantization to bits >= 0.
  """

  def __init__(self,
               bits=8,
               integer=0,
               use_sigmoid=0,
               negative_slope=0.0,
               use_stochastic_rounding=False,
               relu_upper_bound=None,
               is_quantized_clip=True,
               qnoise_factor=1.0,
               var_name=None,
               use_ste=True,
               use_variables=False):
    super(quantized_relu, self).__init__()
    self.bits = bits
    self.integer = integer
    self.use_sigmoid = use_sigmoid
    self.negative_slope = negative_slope
    self.use_stochastic_rounding = use_stochastic_rounding
    self.relu_upper_bound = relu_upper_bound
    self.is_quantized_clip = is_quantized_clip
    self.qnoise_factor = qnoise_factor
    self.use_ste = use_ste
    assert negative_slope >= 0.0
    if negative_slope != 0.0:
      assert np.mod(np.log2(negative_slope), 1) == 0
    self.var_name = var_name
    self.use_variables = use_variables

  def __str__(self):
    # Converts Tensors to printable strings by converting to a numpy array and
    # then using regex to remove brackets when there is only one integer bit
    integer_bits = re.sub(
        r"\[(\d)\]", r"\g<1>",
        str(self.integer.numpy() if isinstance(self.integer, tf.Variable
                                              ) else self.integer))

    flags = [str(self.bits), integer_bits]
    if self.use_sigmoid or self.use_stochastic_rounding:
      flags.append(str(int(self.use_sigmoid)))
    if self.negative_slope:
      flags.append(str(self.negative_slope))
    if self.use_stochastic_rounding:
      flags.append(str(int(self.use_stochastic_rounding)))
    return "quantized_relu(" + ",".join(flags) + ")"

  def __call__(self, x):
    if not self.built:
      self.build(var_name=self.var_name, use_variables=self.use_variables)

    non_sign_bits = self.bits - (self.negative_slope != 0.0)
    x = K.cast(x, dtype="float32")
    m = K.cast(K.pow(2, non_sign_bits), dtype="float32")
    m_i = K.cast(K.pow(2, self.integer), dtype="float32")

    # is_quantized_clip has precedence over relu_upper_bound for backward
    # compatibility.
    m_f = K.cast(
        K.pow(
            tf.constant(2., tf.float32),
            K.cast(self.integer, dtype="float32") - non_sign_bits),
        dtype="float32")
    if self.is_quantized_clip:
      x_u = tf.where(x <= m_i - m_f, K.relu(x, alpha=self.negative_slope),
                     tf.ones_like(x) * (m_i - m_f))
    elif self.relu_upper_bound is not None:
      x_u = tf.where(x <= self.relu_upper_bound,
                     K.relu(x, alpha=self.negative_slope),
                     tf.ones_like(x) * self.relu_upper_bound)
    else:
      x_u = K.relu(x, alpha=self.negative_slope)

    if self.use_sigmoid:
      p = _sigmoid(x / m_i) * m
      xq = m_i * tf.keras.backend.clip(
          2.0 * (_round_through(p, self.use_stochastic_rounding) / m) - 1.0,
          0.0, 1.0 - 1.0 / m)
      if self.negative_slope > 0:
        neg_factor = 1 / (self.negative_slope * m)
        xq = xq + m_i * self.negative_slope * tf.keras.backend.clip(
            2.0 * (_round_through(p * self.negative_slope,
                                  self.use_stochastic_rounding) * neg_factor) -
            1.0, -1.0, 0.0)
    else:
      p = x * m / m_i
      xq = m_i * tf.keras.backend.clip(
          _round_through(p, self.use_stochastic_rounding) / m, 0.0,
          1.0 - 1.0 / m)
      if self.negative_slope > 0:
        neg_factor = 1 / (self.negative_slope * m)
        xq = xq + m_i * self.negative_slope * (
            tf.keras.backend.clip(
                _round_through(p * self.negative_slope,
                               self.use_stochastic_rounding) * neg_factor, -1.0,
                0.0))

    if self.relu_upper_bound and not self.is_quantized_clip:
      xq = tf.where(xq <= self.relu_upper_bound, xq,
                    tf.ones_like(xq) * self.relu_upper_bound)

    if self.use_ste:
      return x_u + tf.stop_gradient(self.qnoise_factor * (-x_u + xq))
    else:
      return (1 - self.qnoise_factor) * x_u + tf.stop_gradient(
          self.qnoise_factor * xq)

  def max(self):
    """Get the maximum value that quantized_relu can represent."""
    unsigned_bits = self.bits - (self.negative_slope != 0.0)

    if unsigned_bits > 0:
      return max(
          1.0,
          np.array(
              K.pow(2.0, K.cast(self.integer, dtype="float32")),
              dtype="float32"))
    else:
      return 1.0

  def min(self):
    """Get the minimum value that quantized_relu can represent."""
    if self.negative_slope == 0.0:
      return 0.0

    unsigned_bits = self.bits - 1
    if unsigned_bits > 0:
      return min(
          -0.0, -self.negative_slope * np.array(
              K.pow(2.0, K.cast(self.integer, dtype="float32")),
              dtype="float32"))
    else:
      return -1.0

  def range(self):
    """Returns a list of all values that quantized_relu can represent

      ordered by their binary representation ascending.
    """
    assert self.use_sigmoid == 0  # current unsupported
    assert self.negative_slope == 0  # # unsupported unsupported
    x = np.asarray(range(2**self.bits))
    return x * np.array(
        K.pow(2.0, -self.bits + K.cast(self.integer, dtype="float32")),
        dtype="float32")

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "bits":
            self.bits,
        "integer":
            self.integer.numpy() if isinstance(self.integer, tf.Variable) else
            self.integer,
        "use_sigmoid":
            self.use_sigmoid,
        "negative_slope":
            self.negative_slope,
        "use_stochastic_rounding":
            self.use_stochastic_rounding,
        "relu_upper_bound":
            self.relu_upper_bound,
        "qnoise_factor":
            self.qnoise_factor.numpy() if isinstance(
                self.qnoise_factor, tf.Variable) else self.qnoise_factor
    }
    return config class quantized_ulaw(BaseQuantizer):  # pylint: disable=invalid-name
  """Computes a u-law quantization.

  Attributes:
    bits: number of bits to perform quantization.
    integer: number of bits to the left of the decimal point.
    symmetric: if true, we will have the same number of values for positive
      and negative numbers.
    u: parameter of u-law

  Returns:
    Function that performs ulaw + quantization to bits in the range -1.0 to 1.0.
  """

  def __init__(self, bits=8, integer=0, symmetric=0, u=255.0):
    super(quantized_ulaw, self).__init__()
    self.bits = bits
    self.integer = integer
    self.symmetric = symmetric
    self.u = u

  def __str__(self):
    flags = [str(self.bits), str(self.integer)]
    if self.symmetric or self.u != 255.0:
      flags.append(str(int(self.symmetric)))
    if self.u != 255.0:
      flags.append(str(self.u))
    return "quantized_ulaw(" + ",".join(flags) + ")"

  def __call__(self, x):
    non_sign_bits = self.bits - 1
    m = pow(2, non_sign_bits)
    m_i = pow(2, self.integer)
    p = _sigmoid(x / m_i) * m
    rp = 2.0 * (_round_through(p) / m) - 1.0
    u_law_p = tf.sign(rp) * tf.keras.backend.log(
        1 + self.u * tf.abs(rp)) / tf.keras.backend.log(1 + self.u)
    xq = m_i * tf.keras.backend.clip(u_law_p, -1.0 +
                                     (1.0 * self.symmetric) / m, 1.0 - 1.0 / m)
    return xq

  def max(self):
    """Get the maximum value that quantized_ulaw can represent."""
    unsigned_bits = self.bits - 1

    if unsigned_bits > 0:
      return max(1.0, np.power(2.0, self.integer))
    else:
      return 1.0

  def min(self):
    """Get the minimum value that quantized_ulaw can represent."""
    unsigned_bits = self.bits - 1

    if unsigned_bits > 0:
      return -max(1.0, np.power(2.0, self.integer))
    else:
      return -1.0

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "bits": self.bits,
        "integer": self.integer,
        "symmetric": self.symmetric,
        "u": self.u
    }
    return config class quantized_tanh(BaseQuantizer):  # pylint: disable=invalid-name
  """Computes a quantized tanh to a number of bits.

  Modified from:

  [https://github.com/BertMoons/QuantizedNeuralNetworks-Keras-Tensorflow]

  Attributes:
    bits: number of bits to perform quantization.
    use_stochastic_rounding: if true, we perform stochastic rounding.
    symmetric: if true, we will have the same number of values for positive
      and negative numbers.
    use_real_tanh: if true, use the tanh function from Keras backend,
      if false, use tanh that is defined as 2 * sigmoid(x) - 1

  Returns:
    Function that performs tanh + quantization to bits in the range -1.0 to 1.0.
  """

  def __init__(self, bits=8, use_stochastic_rounding=False,
               symmetric=False, use_real_tanh=False):
    super(quantized_tanh, self).__init__()
    self.bits = bits
    self.symmetric = symmetric
    self.use_stochastic_rounding = use_stochastic_rounding
    self.use_real_tanh = use_real_tanh

  def __str__(self):
    flags = [str(self.bits)]
    if self.use_stochastic_rounding:
      flags.append(str(int(self.use_stochastic_rounding)))
    if self.symmetric:
      flags.append(str(int(self.symmetric)))
    if self.use_real_tanh:
      flags.append(str(int(self.use_real_tanh)))
    return "quantized_tanh(" + ",".join(flags) + ")"

  def __call__(self, x):
    non_sign_bits = self.bits - 1
    x = K.cast_to_floatx(x)
    m = K.cast_to_floatx(K.pow(2, non_sign_bits))
    p = K.tanh(x) if self.use_real_tanh else 2.0 * _sigmoid(x) - 1.0
    return tf.keras.backend.clip(
                                 (_round_through(p * m, self.use_stochastic_rounding) / m),
                                 -1.0 + (1.0 * self.symmetric) / m,
                                 1.0 - 1.0 / m)

  def max(self):
    """Get the maximum value that quantized_tanh can represent."""
    return 1.0 - 1.0 / pow(2, self.bits - 1)

  def min(self):
    """Get the minimum value that quantized_tanh can represent."""
    return -1.0 + (1.0 * self.symmetric) / pow(2, self.bits - 1)

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "bits": self.bits,
        "symmetric": self.symmetric,
        "use_stochastic_rounding": self.use_stochastic_rounding,
        "use_real_tanh": self.use_real_tanh
    }
    return config class quantized_sigmoid(BaseQuantizer):  # pylint: disable=invalid-name
  """Computes a quantized sigmoid to a number of bits.

  Attributes:
    bits: number of bits to perform quantization.
    symmetric: if true, we will have the same number of values for positive
      and negative numbers.
    use_real_sigmoid: if true, will use the sigmoid from Keras backend
    use_stochastic_rounding: if true, we perform stochastic rounding.

  Returns:
    Function that performs sigmoid + quantization to bits in the range 0.0 to 1.0.
  """

  def __init__(self, bits=8, symmetric=False,
               use_real_sigmoid=False,
               use_stochastic_rounding=False):
    super(quantized_sigmoid, self).__init__()
    self.bits = bits
    self.symmetric = symmetric
    self.use_real_sigmoid = use_real_sigmoid
    self.use_stochastic_rounding = use_stochastic_rounding

  def __str__(self):
    flags = [str(self.bits)]
    if self.symmetric:
      flags.append(str(int(self.symmetric)))
    if self.use_real_sigmoid:
      flags.append(str(int(self.use_real_sigmoid)))
    if self.use_stochastic_rounding:
      flags.append(str(int(self.use_stochastic_rounding)))
    return "quantized_sigmoid(" + ",".join(flags) + ")"

  def __call__(self, x):
    x = K.cast_to_floatx(x)
    m = K.cast_to_floatx(K.pow(2, self.bits))

    p = K.sigmoid(x) if self.use_real_sigmoid else _sigmoid(x)

    return tf.keras.backend.clip((_round_through(p*m, self.use_stochastic_rounding) / m),
                                 (1.0 * self.symmetric) / m,
                                 1.0 - 1.0 / m)

  def max(self):
    """Get the maximum value that quantized_sigmoid can represent."""
    return 1.0 - 1.0 / pow(2, self.bits)

  def min(self):
    """Get the minimum value that quantized_sigmoid can represent."""
    return (1.0 * self.symmetric) / pow(2, self.bits)

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "bits": self.bits,
        "symmetric": self.symmetric,
        "use_real_sigmoid": self.use_real_sigmoid,
        "use_stochastic_rounding": self.use_stochastic_rounding
    }
    return config class quantized_po2(BaseQuantizer):  # pylint: disable=invalid-name
  """Quantizes to the closest power of 2.

  Attributes:
    bits: An integer, the bits allocated for the exponent, its sign and the sign
      of x.
    max_value: An float or None. If None, no max_value is specified.
      Otherwise, the maximum value of quantized_po2 <= max_value
    use_stochastic_rounding: A boolean, default is False, if True, it uses
      stochastic rounding and forces the mean of x to be x statstically.
    quadratic_approximation: A boolean, default is False if True, it forces the
      exponent to be even number that closted to x.
    log2_rounding: A string, log2 rounding mode. "rnd" and "floor" currently
      supported, corresponding to tf.round and tf.floor respectively.
    qnoise_factor: float. a scalar from 0 to 1 that represents the level of
      quantization noise to add. This controls the amount of the quantization
      noise to add to the outputs by changing the weighted sum of
      (1 - qnoise_factor)*unquantized_x + qnoise_factor*quantized_x.
    var_name: String or None. A variable name shared between the tf.Variables
      created in the build function. If None, it is generated automatically.
    use_ste: Bool. Whether to use "straight-through estimator" (STE) method or
        not.
    use_variables: Bool. Whether to make the quantizer variables to be dynamic
      tf.Variables or not.
  """

  def __init__(self,
               bits=8,
               max_value=None,
               use_stochastic_rounding=False,
               quadratic_approximation=False,
               log2_rounding="rnd",
               qnoise_factor=1.0,
               var_name=None,
               use_ste=True,
               use_variables=False):
    super(quantized_po2, self).__init__()
    self.bits = bits
    self.max_value = max_value
    self.use_stochastic_rounding = use_stochastic_rounding
    self.log2_rounding = log2_rounding
    # if True, round to the exponent for sqrt(x),
    # so that the return value can be divided by two without remainder.
    self.quadratic_approximation = quadratic_approximation
    need_exponent_sign_bit = _need_exponent_sign_bit_check(self.max_value)
    non_sign_bits = self.bits - 1
    self._min_exp, self._max_exp = _get_min_max_exponents(
        non_sign_bits, need_exponent_sign_bit, self.quadratic_approximation)
    # qnoise_factor related attributes
    self.qnoise_factor = qnoise_factor
    self.use_ste = use_ste
    self.var_name = var_name
    self.use_variables = use_variables

  def __str__(self):
    flags = [str(self.bits)]
    if self.max_value is not None or self.use_stochastic_rounding:
      flags.append(str(int(self.max_value)))
    if self.use_stochastic_rounding:
      flags.append(str(int(self.use_stochastic_rounding)))
    if self.quadratic_approximation:
      flags.append(
          "quadratic_approximation=" + str(int(self.quadratic_approximation)))
    return "quantized_po2(" + ",".join(flags) + ")"

  def __call__(self, x):
    if not self.built:
      self.build(var_name=self.var_name, use_variables=self.use_variables)

    x_sign = tf.sign(x)
    x_sign += (1.0 - tf.abs(x_sign))
    x_abs = tf.abs(x)
    x_clipped = _clip_power_of_two(x_abs, self._min_exp, self._max_exp,
                                   self.max_value,
                                   self.quadratic_approximation,
                                   self.use_stochastic_rounding,
                                   self.log2_rounding)
    xq = x_sign * pow(2.0, x_clipped)

    if self.use_ste:
      return x + tf.stop_gradient(self.qnoise_factor * (-x + xq))
    else:
      return (1 - self.qnoise_factor) * x + tf.stop_gradient(
          self.qnoise_factor * xq)

  def max(self):
    """Get the maximum value that quantized_po2 can represent."""
    if self.max_value:
      return max(1.0, self.max_value)
    else:
      return max(1.0, 2**self._max_exp)

  def min(self):
    """Get the minimum value that quantized_po2 can represent."""
    if self.max_value:
      return -max(1.0, self.max_value)
    else:
      return -max(1.0, 2**self._max_exp)

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    """Gets configugration of the quantizer.

    Returns:
      A dict mapping quantization configuration, including
        bits: bitwidth for exponents.
        max_value: the maximum value of this quantized_po2 can represent.
        use_stochastic_rounding:
          if True, stochastic rounding is used.
        quadratic_approximation:
          if True, the exponent is enforced to be even number, which is
          the closest one to x.
        log2_rounding:
          A string, Log2 rounding mode
    """
    config = {
        "bits":
            self.bits,
        "max_value":
            self.max_value,
        "use_stochastic_rounding":
            self.use_stochastic_rounding,
        "quadratic_approximation":
            self.quadratic_approximation,
        "qnoise_factor":
            self.qnoise_factor.numpy() if isinstance(
                self.qnoise_factor, tf.Variable) else self.qnoise_factor,
        "log2_rounding":
            self.log2_rounding
    }
    return config class quantized_relu_po2(BaseQuantizer):  # pylint: disable=invalid-name
  """Quantizes x to the closest power of 2 when x > 0

  Attributes:
    bits: An integer, the bits allocated for the exponent and its sign.
    max_value: default is None, or a non-negative value to put a constraint for
      the max value.
    negative_slope: slope when activation < 0, needs to be power of 2.
    use_stochastic_rounding: A boolean, default is False, if True, it uses
      stochastic rounding and forces the mean of x to be x statstically.
    quadratic_approximation: A boolean, default is False if True, it forces the
      exponent to be even number that is closest to x.
    log2_rounding: A string, log2 rounding mode. "rnd" and "floor" currently
      supported, corresponding to tf.round and tf.floor respectively.
    qnoise_factor: float. a scalar from 0 to 1 that represents the level of
      quantization noise to add. This controls the amount of the quantization
      noise to add to the outputs by changing the weighted sum of
      (1 - qnoise_factor)*unquantized_x + qnoise_factor*quantized_x.
    var_name: String or None. A variable name shared between the tf.Variables
      created in the build function. If None, it is generated automatically.
    use_ste: Bool. Whether to use "straight-through estimator" (STE) method or
        not.
    use_variables: Bool. Whether to make the quantizer variables to be dynamic
      tf.Variables or not.
  """

  def __init__(self,
               bits=8,
               max_value=None,
               negative_slope=0,
               use_stochastic_rounding=False,
               quadratic_approximation=False,
               log2_rounding="rnd",
               qnoise_factor=1.0,
               var_name=None,
               use_ste=True,
               use_variables=False):
    super(quantized_relu_po2, self).__init__()
    self.bits = bits
    self.max_value = max_value
    self.negative_slope = negative_slope
    self.use_stochastic_rounding = use_stochastic_rounding
    self.log2_rounding = log2_rounding
    # if True, round to the exponent for sqrt(x),
    # so that the return value can be divided by two without remainder.
    self.quadratic_approximation = quadratic_approximation
    need_exponent_sign_bit = _need_exponent_sign_bit_check(self.max_value)
    self._min_exp = -2**(self.bits - need_exponent_sign_bit)
    self._max_exp = 2**(self.bits - need_exponent_sign_bit) - 1
    if self.quadratic_approximation:
      self._max_exp = 2 * (self._max_exp // 2)

    assert negative_slope >= 0.0
    if negative_slope != 0:
      assert np.mod(np.log2(negative_slope), 1) == 0
    # qnoise_factor related attributes
    self.qnoise_factor = qnoise_factor
    self.use_ste = use_ste
    self.var_name = var_name
    self.use_variables = use_variables

  def __str__(self):
    flags = [str(self.bits)]
    if self.max_value is not None or self.use_stochastic_rounding:
      flags.append(str(int(self.max_value)))
    if self.negative_slope:
      flags.append(str(self.negative_slope))
    if self.use_stochastic_rounding:
      flags.append(str(int(self.use_stochastic_rounding)))
    if self.quadratic_approximation:
      flags.append(
          "quadratic_approximation=" + str(int(self.quadratic_approximation)))
    return "quantized_relu_po2(" + ",".join(flags) + ")"

  def __call__(self, x):
    if not self.built:
      self.build(var_name=self.var_name, use_variables=self.use_variables)

    x_original = x

    if self.max_value is None:
      x = K.relu(x, self.negative_slope)
    else:
      x = tf.where(
          x <= self.max_value,
          K.relu(x, self.negative_slope),
          tf.ones_like(x) * self.max_value)

    x_pos_clipped = _clip_power_of_two(
        K.relu(x_original),
        self._min_exp, self._max_exp,
        self.max_value,
        self.quadratic_approximation,
        self.use_stochastic_rounding,
        self.log2_rounding)

    x_neg_clipped = _clip_power_of_two(
        K.relu(-x_original) * self.negative_slope,
        self._min_exp, self._max_exp,
        self.max_value,
        self.quadratic_approximation,
        self.use_stochastic_rounding,
        self.log2_rounding)

    xq = tf.where(
        tf.logical_or(x_original >= 0.0, self.negative_slope == 0.0),
        pow(2.0, x_pos_clipped), -pow(2.0, x_neg_clipped))

    if self.use_ste:
      return x + tf.stop_gradient(self.qnoise_factor * (-x + xq))
    else:
      return (1 - self.qnoise_factor) * x + tf.stop_gradient(
          self.qnoise_factor * xq)

  def max(self):
    """Get the maximum value that quantized_relu_po2 can represent."""
    if self.max_value:
      return max(1.0, self.max_value)
    else:
      return max(1.0, 2**self._max_exp)

  def min(self):
    """Get the minimum value that quantized_relu_po2 can represent."""
    if self.negative_slope == 0.0:
      return 2**self._min_exp

    unsigned_bits = self.bits - 1
    if unsigned_bits > 0:
      return min(2**self._min_exp, - self.negative_slope * np.power(2.0, unsigned_bits))
    else:
      return 2**self._min_exp

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    """Gets configugration of the quantizer.

    Returns:
      A dict mapping quantization configuration, including
        bits: bitwidth for exponents.
        max_value: the maximum value of this quantized_relu_po2 can represent.
        use_stochastic_rounding:
          if True, stochastic rounding is used.
        quadratic_approximation:
          if True, the exponent is enforced to be even number, which is
          the closest one to x.
        log2_rounding:
          A string, Log2 rounding mode

    """

    config = {
        "bits":
            self.bits,
        "max_value":
            self.max_value,
        "negative_slope":
            self.negative_slope,
        "use_stochastic_rounding":
            self.use_stochastic_rounding,
        "quadratic_approximation":
            self.quadratic_approximation,
        "qnoise_factor":
            self.qnoise_factor.numpy() if isinstance(
                self.qnoise_factor, tf.Variable) else self.qnoise_factor,
        "log2_rounding":
            self.log2_rounding
    }
    return config class stochastic_binary(binary):  # pylint: disable=invalid-name
  """Computes a stochastic activation function returning -alpha or +alpha.

  Computes straight-through approximation using random sampling to make
  E[dL/dy] = E[dL/dx], and computing the sign function. See explanation above.

  Attributes:
    x: tensor to perform sign opertion with stochastic sampling.
    alpha: binary is -alpha or +alpha, or "auto" or "auto_po2".
    bits: number of bits to perform quantization.
    temperature: amplifier factor for sigmoid function, making stochastic
        behavior less stochastic as it moves away from 0.
    use_real_sigmoid: use real sigmoid from tensorflow for probablity.

  Returns:
    Computation of sign with stochastic sampling with straight through gradient.
  """

  def __init__(self, alpha=None, temperature=6.0, use_real_sigmoid=True):
    super(stochastic_binary, self).__init__(alpha=alpha)
    self.alpha = alpha
    self.bits = 1
    self.temperature = temperature
    self.use_real_sigmoid = use_real_sigmoid
    self.default_alpha = 1.0
    self.scale = None

  def __str__(self):
    flags = []
    if self.alpha is not None:
      alpha = str(self.alpha)
      if isinstance(self.alpha, six.string_types):
        alpha = "'" + alpha + "'"
      flags.append("alpha=" + alpha)
    if self.temperature != 6.0:
      flags.append("temperature=" + str(self.temperature))
    if not self.use_real_sigmoid:
      flags.append("use_real_sigmoid=" + str(int(self.use_real_sigmoid)))
    return "stochastic_binary(" + ",".join(flags) + ")"

  def __call__(self, x):
    def stochastic_output():
      if isinstance(self.alpha, six.string_types):
        assert self.alpha in ["auto", "auto_po2"]
        len_axis = len(x.shape)
        if len_axis > 1:
          if K.image_data_format() == "channels_last":
            axis = list(range(len_axis - 1))
          else:
            axis = list(range(1, len_axis))
        else:
          axis = [0]
        std = K.std(x, axis=axis, keepdims=True) + K.epsilon()
      else:
        std = 1.0

      if self.use_real_sigmoid:
        p = tf.keras.backend.sigmoid(self.temperature * x / std)
      else:
        p = _sigmoid(self.temperature * x / std)

      r = tf.random.uniform(tf.shape(x))
      q = tf.sign(p - r)
      q += (1.0 - tf.abs(q))
      q_non_stochastic = tf.sign(x)
      q_non_stochastic += (1.0 - tf.abs(q_non_stochastic))
      scale = _get_scale(self.alpha, x, q_non_stochastic)
      self.scale = scale
      return x + tf.stop_gradient(-x + scale * q)

    output = tf_utils.smart_cond(
        K.learning_phase(), stochastic_output, lambda: binary.__call__(self, x))
    return output

  def _set_trainable_parameter(self):
    if self.alpha is None:
      self.alpha = "auto_po2"

  def max(self):
    """Get the maximum value that stochastic_binary can respresent."""
    if self.alpha is None or isinstance(self.alpha, six.string_types):
      return 1.0
    else:
      return max(1.0, self.alpha)

  def min(self):
    """Get the minimum value that stochastic_binary can respresent."""
    if self.alpha is None or isinstance(self.alpha, six.string_types):
      return -1.0
    else:
      return -max(1.0, self.alpha)

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "alpha": self.alpha,
        "temperature": self.temperature,
        "use_real_sigmoid": self.use_real_sigmoid,
    }
    return config class stochastic_ternary(ternary):  # pylint: disable=invalid-name
  """Computes a stochastic activation function returning -alpha, 0 or +alpha.

  Computes straight-through approximation using random sampling to make
  E[dL/dy] = E[dL/dx], and computing the sign function. See explanation above.

  Attributes:
    x: tensor to perform sign opertion with stochastic sampling.
    bits: number of bits to perform quantization.
    alpha: ternary is -alpha or +alpha, or "auto" or "auto_po2".
    threshold: (1-threshold) specifies the spread of the +1 and -1 values.
    temperature: amplifier factor for sigmoid function, making stochastic
      less stochastic as it moves away from 0.
    use_real_sigmoid: use real sigmoid for probability.
    number_of_unrolls: number of times we iterate between scale and threshold.

  Returns:
    Computation of sign with stochastic sampling with straight through gradient.
  """

  def __init__(self, alpha=None, threshold=None, temperature=8.0,
               use_real_sigmoid=True, number_of_unrolls=5):
    super(stochastic_ternary, self).__init__(
      alpha=alpha,
      threshold=threshold,
      number_of_unrolls=number_of_unrolls)

    self.bits = 2
    self.alpha = alpha
    self.threshold = threshold
    assert threshold != 1.0
    self.default_alpha = 1.0
    self.default_threshold = 0.33
    self.temperature = temperature
    self.use_real_sigmoid = use_real_sigmoid
    self.number_of_unrolls = number_of_unrolls
    self.scale = None

  def __str__(self):
    flags = []
    if self.alpha is not None:
      alpha = str(self.alpha)
      if isinstance(self.alpha, six.string_types):
        alpha = "'" + alpha + "'"
      flags.append("alpha=" + alpha)
    if self.threshold is not None:
      flags.append("threshold=" + str(self.threshold))
    if self.temperature != 8.0:
      flags.append("temperature=" + str(self.temperature))
    if not self.use_real_sigmoid:
      flags.append("use_real_sigmoid=0")
    if self.number_of_unrolls != 5:
      flags.append("number_of_unrolls=" + str(self.number_of_unrolls))
    return "stochastic_ternary(" + ",".join(flags) + ")"

  def __call__(self, x):
    def stochastic_output():
      # right now we only accept alpha = "auto" or "auto_po2"

      assert isinstance(self.alpha, six.string_types)
      assert self.alpha in ["auto", "auto_po2"]

      if self.alpha is None:
        scale = self.default_alpha
      elif isinstance(self.alpha, six.string_types):
        scale = 1.0
        assert self.alpha in ["auto", "auto_po2"]
      else:
        assert self.alpha >= 0.0
        scale = float(self.alpha)

      len_axis = len(x.shape)
      if len_axis > 1:
        if K.image_data_format() == "channels_last":
          axis = list(range(len_axis - 1))
        else:
          axis = list(range(1, len_axis))
      else:
        axis = [0]

      x_std = K.std(x, axis=axis, keepdims=True)

      m = K.max(tf.abs(x), axis=axis, keepdims=True)
      scale = 2.*m/3.
      if self.alpha == "auto_po2":
        scale = K.pow(2.0,
                      tf.math.round(K.log(scale + K.epsilon()) / np.log(2.0)))
      for _ in range(self.number_of_unrolls):
        T = scale / 2.0
        q_ns = K.cast(tf.abs(x) >= T, K.floatx()) * K.sign(x)
        scale = _get_scale(self.alpha, x, q_ns)

      x_norm = x / (x_std + K.epsilon())
      T = scale / (2.0 * (x_std + K.epsilon()))

      if self.use_real_sigmoid:
        p0 = tf.keras.backend.sigmoid(self.temperature * (x_norm - T))
        p1 = tf.keras.backend.sigmoid(self.temperature * (x_norm + T))
      else:
        p0 = _sigmoid(self.temperature * (x_norm - T))
        p1 = _sigmoid(self.temperature * (x_norm + T))
      r0 = tf.random.uniform(tf.shape(p0))
      r1 = tf.random.uniform(tf.shape(p1))
      q0 = tf.sign(p0 - r0)
      q0 += (1.0 - tf.abs(q0))
      q1 = tf.sign(p1 - r1)
      q1 += (1.0 - tf.abs(q1))

      q = (q0 + q1) / 2.0
      self.scale = scale
      return x + tf.stop_gradient(-x + scale * q)

    output = tf_utils.smart_cond(
        K.learning_phase(),
        stochastic_output,
        lambda: ternary.__call__(self, x))
    return output

  def _set_trainable_parameter(self):
    if self.alpha is None:
      self.alpha = "auto_po2"

  def max(self):
    """Get the maximum value that stochastic_ternary can respresent."""
    if self.alpha is None or isinstance(self.alpha, six.string_types):
      return 1.0
    else:
      return max(1.0, self.alpha)

  def min(self):
    """Get the minimum value that stochastic_ternary can respresent."""
    if self.alpha is None or isinstance(self.alpha, six.string_types):
      return -1.0
    else:
      return -max(1.0, self.alpha)

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "alpha": self.alpha,
        "threshold": self.threshold,
        "temperature": self.temperature,
        "use_real_sigmoid": self.use_real_sigmoid,
        "number_of_unrolls": self.number_of_unrolls
    }
    return config class ternary(BaseQuantizer):  # pylint: disable=invalid-name
  """Computes an activation function returning -alpha, 0 or +alpha.

  Right now we assume two type of behavior. For parameters, we should
  have alpha, threshold and stochastic rounding on. For activations,
  alpha and threshold should be floating point numbers, and stochastic
  rounding should be off.

  Attributes:
    x: tensor to perform sign opertion with stochastic sampling.
    bits: number of bits to perform quantization.
    alpha: ternary is -alpha or +alpha. Alpha can be "auto" or "auto_po2".
    threshold: threshold to apply "dropout" or dead band (0 value). If "auto"
      is specified, we will compute it per output layer.
    use_stochastic_rounding: if true, we perform stochastic rounding.

  Returns:
    Computation of sign within the threshold.
  """

  def __init__(self, alpha=None, threshold=None, use_stochastic_rounding=False,
               number_of_unrolls=5):
    super(ternary, self).__init__()
    self.bits = 2
    self.alpha = alpha
    self.threshold = threshold
    self.use_stochastic_rounding = use_stochastic_rounding
    self.default_alpha = 1.0
    self.default_threshold = 0.33
    self.number_of_unrolls = number_of_unrolls
    self.scale = None

  def __str__(self):
    flags = []
    if self.alpha is not None:
      alpha = str(self.alpha)
      if isinstance(self.alpha, six.string_types):
        alpha = "'" + alpha + "'"
      flags.append("alpha=" + alpha)
    if self.threshold is not None:
      flags.append("threshold=" + str(self.threshold))
    if self.use_stochastic_rounding:
      flags.append(
          "use_stochastic_rounding=" + str(int(self.use_stochastic_rounding)))
    if self.number_of_unrolls != 5:
      flags.append(
          "number_of_unrolls=" + str(int(self.number_of_unrolls)))
    return "ternary(" + ",".join(flags) + ")"

  def __call__(self, x):
    if isinstance(self.alpha, six.string_types):
      # parameters
      assert self.alpha in ["auto", "auto_po2"]
      assert self.threshold is None
    else:
      # activations
      assert not self.use_stochastic_rounding
      assert not isinstance(self.threshold, six.string_types)

    if self.alpha is None or isinstance(self.alpha, six.string_types):
      scale = 1.0
    elif isinstance(self.alpha, np.ndarray):
      scale = self.alpha
    else:
      scale = float(self.alpha)

    # This is an approximiation from https://arxiv.org/abs/1605.04711
    # We consider channels_last only for now.
    if isinstance(self.alpha, six.string_types):
      # It is for parameters
      # first, compute which asix corresponds to the channels.
      # TODO(hzhuang): support channels_first
      try:
        len_axis = len(x.shape.as_list())
      except AttributeError:
        len_axis = len(list(x.shape))

      if len_axis == 1:
        axis = None
      elif K.image_data_format() == "channels_last":
        axis = list(range(len_axis - 1))
      else:
        axis = list(range(1, len_axis))

      # This approximation is exact if x ~ U[-m, m]. For x ~ N(0, m)
      # we need to iterate a few times before we can coverge
      m = K.max(tf.abs(x), axis=axis, keepdims=True)
      scale = 2 * m / 3.0
      if "po2" in self.alpha:
        scale = K.pow(2.0,
                      tf.math.round(K.log(scale + K.epsilon()) / np.log(2.0)))

      for _ in range(self.number_of_unrolls):
        thres = scale / 2.0
        # once we scale the number precision == 0.33 works
        # well for Uniform and Normal distribution of input
        v = scale * _round_through(
            x / scale,
            use_stochastic_rounding=self.use_stochastic_rounding,
            precision=1. / 3.)
        q = K.cast(tf.abs(v) >= thres, K.floatx()) * tf.sign(x)
        scale = _get_scale(self.alpha, x, q)
    else:
      if self.threshold is None:
        thres = self.default_threshold
      else:
        thres = self.threshold
      q = K.cast(tf.abs(x) >= thres, K.floatx()) * tf.sign(x)

    # ternary ranges from -1 to +1, so we use tanh(x) to be a differentiable
    # version of that.
    if self.alpha is None:
      x = K.tanh(x)

    self.scale = scale
    return x + tf.stop_gradient(-x + scale * q)

  def _set_trainable_parameter(self):
    if self.alpha is None:
      self.alpha = "auto_po2"

  def max(self):
    """Get the maximum value that ternary can respresent."""
    if self.alpha is None or isinstance(self.alpha, six.string_types):
      return 1.0
    else:
      return max(1.0, self.alpha)

  def min(self):
    """Get the minimum value that ternary can respresent."""
    if self.alpha is None or isinstance(self.alpha, six.string_types):
      return -1.0
    else:
      return -max(1.0, self.alpha)

  @classmethod
  def from_config(cls, config):
    return cls(**config)

  def get_config(self):
    config = {
        "alpha": self.alpha,
        "threshold": self.threshold,
        "use_stochastic_rounding": self.use_stochastic_rounding,
        "number_of_unrolls": self.number_of_unrolls
    }
    return config def safe_eval(eval_str, op_dict, *params, **kwparams):  # pylint: disable=invalid-name
  """Replaces eval by a safe eval mechanism."""

  function_split = eval_str.split("(")
  quantizer = op_dict.get(function_split[0], None)

  if len(function_split) == 2:
    args, kwargs = GetParams("(" + function_split[1])
  else:
    args = []
    kwargs = {}

  args = args + list(params)
  for k in kwparams:
    kwargs[k] = kwparams[k]

  # must be Keras activation object if None
  if quantizer is None:
    logging.info("keras dict %s", function_split[0])
    quantizer = keras.activations.get(function_split[0])

  if len(function_split) == 2 or args or kwargs:
    return quantizer(*args, **kwargs)
  else:
    if isinstance(quantizer, type):
      # Check if quantizer is a class
      return quantizer()
    else:
      # Otherwise it is a function, so just return it
      return quantizer
import copy
import json
import tempfile
import types
import numpy as np
import os
import six
import re
import networkx as nx
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras import initializers
from tensorflow.python.keras.layers.core import TFOpLambda
from tensorflow.keras.models import Model
from tensorflow.keras.models import model_from_json
from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper
from tensorflow_model_optimization.python.core.sparsity.keras import prune_registry
from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
from .qlayers import Clip
from .qconv2d_batchnorm import QConv2DBatchnorm
from .qdepthwiseconv2d_batchnorm import QDepthwiseConv2DBatchnorm
from .qlayers import QActivation
from .qlayers import QAdaptiveActivation
from .qpooling import QAveragePooling2D
from .qlayers import QDense
from .qlayers import QInitializer
from .qconvolutional import QConv1D
from .qconvolutional import QConv2D
from .qconvolutional import QConv2DTranspose
from .qrecurrent import QSimpleRNN
from .qrecurrent import QSimpleRNNCell
from .qrecurrent import QLSTM
from .qrecurrent import QLSTMCell
from .qrecurrent import QGRU
from .qrecurrent import QGRUCell
from .qrecurrent import QBidirectional
from .qconvolutional import QSeparableConv1D
from .qconvolutional import QSeparableConv2D
from .qconvolutional import QDepthwiseConv2D
from .qnormalization import QBatchNormalization
from .qpooling import QGlobalAveragePooling2D
from .qtools import qgraph
from .quantizers import binary
from .quantizers import bernoulli
from .quantizers import get_weight_scale
from .quantizers import quantized_bits
from .quantizers import quantized_relu
from .quantizers import quantized_ulaw
from .quantizers import quantized_tanh
from .quantizers import quantized_sigmoid
from .quantizers import quantized_po2
from .quantizers import quantized_relu_po2
from .quantizers import stochastic_binary
from .quantizers import stochastic_ternary
from .quantizers import ternary
from .safe_eval import safe_eval
from tensorflow.python.ops import math_ops
      # This is to avoid unwanted transformations.
      if quantizer is None:
        continue

      if layer["class_name"] == "LeakyReLU":
        negative_slope = layer["config"]["alpha"]
      elif layer["class_name"] == "relu":
        max_value = layer["config"]["max_value"]
        negative_slope = layer["config"]["alpha"]
        threshold = layer["config"]["threshold"]
      else: # ReLU from mobilenet
        max_value = layer["config"]["max_value"]
        negative_slope = layer["config"]["negative_slope"]
        threshold = layer["config"]["threshold"]

      if negative_slope > 0:
        q_name = "leakyrelu"
      else:
        q_name = "relu"

      # If quantizer exists in dictionary related to this name,
      # use it, otherwise, use normal transformations.

      if not isinstance(quantizer, dict) or quantizer.get(q_name, None):
        # Only change activation layer if we will use a quantized activation.

        layer["class_name"] = "QActivation"

        # Remove relu specific configurations
        # remember that quantized relu's are always upper bounded.

        if layer["class_name"] == "LeakyReLU":
          del layer["config"]["alpha"]
        elif layer["class_name"] == "relu":
          del layer["config"]["max_value"]
          del layer["config"]["alpha"]
          del layer["config"]["threshold"]
        else: # ReLU from mobilenet
          del layer["config"]["max_value"]
          del layer["config"]["negative_slope"]
          del layer["config"]["threshold"]

        if isinstance(quantizer, dict):
          quantizer = quantizer[q_name]
        if quantizer:
          layer["config"]["activation"] = quantizer
        else:
          quantize_activation(layer["config"], activation_bits)

    elif layer["class_name"] == "BatchNormalization":
      # We will assume at least QBatchNormalization or
      # layer name is in dictionary to enable conversion
      # otherwise we will just skip it.
      if (
          layer_config["name"] not in quantizer_config and
          "QBatchNormalization" not in quantizer_config
      ):
        continue

      layer["class_name"] = "QBatchNormalization"
      # Needs to add kernel/bias quantizers.
      gamma_quantizer = get_config(
          quantizer_config, layer, "QBatchNormalization",
          "gamma_quantizer")
      beta_quantizer = get_config(
          quantizer_config, layer, "QBatchNormalization",
          "beta_quantizer")
      mean_quantizer = get_config(
          quantizer_config, layer, "QBatchNormalization",
          "mean_quantizer")
      variance_quantizer = get_config(
          quantizer_config, layer, "QBatchNormalization",
          "variance_quantizer")

      layer_config["gamma_quantizer"] = gamma_quantizer
      layer_config["beta_quantizer"] = beta_quantizer
      layer_config["mean_quantizer"] = mean_quantizer
      layer_config["variance_quantizer"] = variance_quantizer

    elif layer["class_name"] in ["AveragePooling2D", "GlobalAveragePooling2D"]:
      q_name = "Q" + layer["class_name"]
      # Adds the average quanizer to config.
      average_quantizer = get_config(
          quantizer_config, layer, q_name, "average_quantizer")

      # This is to avoid unwanted transformations.
      if average_quantizer is None:
        continue

      layer["class_name"] = q_name

      layer_config["average_quantizer"] = average_quantizer

      # Adds activation to config.
      quantizer = get_config(
          quantizer_config, layer, q_name, "activation_quantizer")

      if quantizer:
        layer_config["activation"] = quantizer
      else:
        quantize_activation(layer_config, activation_bits)

  # We need to keep a dictionary of custom objects as our quantized library
  # is not recognized by keras.

  qmodel = quantized_model_from_json(json.dumps(jm), custom_objects)

  # If transfer_weights is true, we load the weights from model to qmodel.

  if transfer_weights and not enable_bn_folding:
    for layer, qlayer in zip(model.layers, qmodel.layers):
      if layer.get_weights():
        qlayer.set_weights(copy.deepcopy(layer.get_weights()))

  return qmodel


def _add_supported_quantized_objects(custom_objects):
  """Map all the quantized objects."""
  custom_objects["QInitializer"] = QInitializer

  custom_objects["QDense"] = QDense