Repository: larq/larq

Issue:
Documentation (Regarding to JOSS review)
Some suggestions:
In the quantizers documentation:
1. It seems that SteSign/ste_sign (if they are the same) is repeated twice [here](https://larq.dev/api/quantizers/#ste_sign). The same for swish_sign/SwishSign.
2. These quantizers (magnitude_aware_sign,swish_sign)have no formula, nor description. 





Discussion:
Thanks for the suggestions. The feedback is really valuable for us 👍 

> 1. It seems that SteSign/ste_sign (if they are the same) is repeated twice here. The same for swish_sign/SwishSign.

Indeed they are computing the same thing. The difference is the usage. Larq allows three different usage patterns outlined [here](https://larq.dev/api/quantizers/#larqquantizers):
1. In a larq layer the quantizer can be referenced with a string. Compare `kernel_quantizer="ste_sign"` and `activation="relu"`.
2. The quantizer can also be called as a function. Compare `kernel_quantizer=lq.quantizers.ste_sign` and `activation=tf.nn.relu`.
3. If the default arguments of the quantizer need to be configured, it is important that the quantizer class is serializable. That means the quantizer should be called like `kernel_quantizer=lq.quantizers.SteSign(clip_value=1.25)` similar to `kernel_constraint=tf.keras.constraints.MaxNorm(max_value=2)`

Options 1 and 2 are really only here for ease of use. What we could do though is remove option 2 and replace it with `kernel_quantizer=lq.quantizers.ste_sign()` which would than mimik [Keras Constraints](https://keras.io/constraints/) and would only require us to alias `ste_sign` to `SteSign` so both would be classes which would make implementation simpler and usage consistent as well.

Though I am not sure if people would expect quantizers to be used like activations.
@EduPH @koenhelwegen @timdebruin @AdamHillier What do you think?

> 2. These quantizers (magnitude_aware_sign,swish_sign)have no formula, nor description.

I added more docs and formulas in #257
I think that it is natural to use them as activation functions. However, if you prefer to keep as it is now, at least it should be clearer on the documentation and unified in some sense. Maybe you could combine both definitions and include the different ways to use them together. I haven't got strong opposition to any option you decide to choose.
> I think that it is natural to use them as activation functions. However, if you prefer to keep as it is now, at least it should be clearer on the documentation and unified in some sense. Maybe you could combine both definitions and include the different ways to use them together. I haven't got strong opposition to any option you decide to choose.

I have to think about this a bit more, of what makes sense. For now I reordered the API docs to make the distinction clearer: #274 

Resolving patch:
diff --git a/larq/conftest.py b/larq/conftest.py
--- a/larq/conftest.py
+++ b/larq/conftest.py
@@ -19,6 +19,7 @@ def graph_mode():
     with context.graph_mode():
         with tf.compat.v1.Session().as_default():
             yield
+            tf.keras.backend.clear_session()
 
 
 @pytest.fixture(params=["eager", "graph"])
@@ -28,6 +29,7 @@ def eager_and_graph_mode(request):
         with context.graph_mode():
             with tf.compat.v1.Session().as_default():
                 yield request.param
+                tf.keras.backend.clear_session()
     else:
         with context.eager_mode():
             yield request.param
diff --git a/larq/quantizers.py b/larq/quantizers.py
--- a/larq/quantizers.py
+++ b/larq/quantizers.py
@@ -41,9 +41,6 @@
 lq.layers.QuantDense(64, kernel_quantizer="ste_sign")
 ```
 ```python
-lq.layers.QuantDense(64, kernel_quantizer=lq.quantizers.ste_sign)
-```
-```python
 lq.layers.QuantDense(64, kernel_quantizer=lq.quantizers.SteSign(clip_value=1.0))
 ```
 """
@@ -55,14 +52,8 @@
 from larq import math, utils
 
 __all__ = [
-    "ste_sign",
-    "approx_sign",
-    "magnitude_aware_sign",
-    "swish_sign",
-    "ste_tern",
-    "ste_heaviside",
-    "dorefa_quantizer",
     "SteSign",
+    "ApproxSign",
     "MagnitudeAwareSign",
     "SwishSign",
     "SteTern",
@@ -80,75 +71,7 @@ def _clipped_gradient(x, dy, clip_value):
     return tf.where(mask, dy, zeros)
 
 
-class QuantizerFunctionWrapper:
-    """Wraps a quantizer function in a class that can be serialized.
-
-    # Arguments
-    fn: The quantizer function to wrap, with signature `fn(x, **kwargs)`.
-    **kwargs: The keyword arguments that are passed on to `fn`.
-    """
-
-    def __init__(self, fn: Callable[[tf.Tensor], tf.Tensor], **kwargs):
-        self.fn = fn
-        self.precision = getattr(fn, "precision", 32)
-        self._fn_kwargs = kwargs
-
-    def __call__(self, x: tf.Tensor) -> tf.Tensor:
-        """Invokes the `QuantizerFunctionWrapper` instance.
-
-        # Arguments
-        x: Input tensor.
-
-        # Returns
-        Quantized tensor.
-        """
-        return self.fn(x, **self._fn_kwargs)
-
-    def get_config(self):
-        return {
-            k: tf.keras.backend.eval(v)
-            if tf.is_tensor(v) or isinstance(v, tf.Variable)
-            else v
-            for k, v in self._fn_kwargs.items()
-        }
-
-
-@utils.register_keras_custom_object
-@utils.set_precision(1)
 def ste_sign(x: tf.Tensor, clip_value: float = 1.0) -> tf.Tensor:
-    r"""Sign binarization function.
-
-    \\[
-    q(x) = \begin{cases}
-      -1 & x < 0 \\\
-      1 & x \geq 0
-    \end{cases}
-    \\]
-
-    The gradient is estimated using the Straight-Through Estimator
-    (essentially the binarization is replaced by a clipped identity on the
-    backward pass).
-    \\[\frac{\partial q(x)}{\partial x} = \begin{cases}
-      1 & \left|x\right| \leq \texttt{clip_value} \\\
-      0 & \left|x\right| > \texttt{clip_value}
-    \end{cases}\\]
-
-    ```plot-activation
-    quantizers.ste_sign
-    ```
-
-    # Arguments
-    x: Input tensor.
-    clip_value: Threshold for clipping gradients. If `None` gradients are not clipped.
-
-    # Returns
-    Binarized tensor.
-
-    # References
-    - [Binarized Neural Networks: Training Deep Neural Networks with Weights and
-      Activations Constrained to +1 or -1](http://arxiv.org/abs/1602.02830)
-    """
-
     @tf.custom_gradient
     def _call(x):
         def grad(dy):
@@ -163,72 +86,8 @@ def _scaled_sign(x):  # pragma: no cover
     return 1.3 * ste_sign(x)
 
 
-@utils.register_keras_custom_object
-@utils.set_precision(1)
-def magnitude_aware_sign(x: tf.Tensor, clip_value: float = 1.0) -> tf.Tensor:
-    r"""Magnitude-aware sign for Bi-Real Net.
-
-    A scaled sign function computed according to Section 3.3 in
-    [Zechun Liu et al](https://arxiv.org/abs/1808.00278).
-
-    ```plot-activation
-    quantizers._scaled_sign
-    ```
-
-    # Arguments
-    x: Input tensor
-    clip_value: Threshold for clipping gradients. If `None` gradients are not clipped.
-
-    # Returns
-    Scaled binarized tensor (with values in \\(\\{-a, a\\}\\), where \\(a\\) is a float).
-
-    # References
-    - [Bi-Real Net: Enhancing the Performance of 1-bit CNNs With Improved
-      Representational Capability and Advanced Training
-      Algorithm](https://arxiv.org/abs/1808.00278)
-
-    """
-    scale_factor = tf.reduce_mean(tf.abs(x), axis=list(range(len(x.shape) - 1)))
-
-    return tf.stop_gradient(scale_factor) * ste_sign(x, clip_value=clip_value)
-
-
-@utils.register_keras_custom_object
-@utils.set_precision(1)
 @tf.custom_gradient
 def approx_sign(x: tf.Tensor) -> tf.Tensor:
-    r"""
-    Sign binarization function.
-    \\[
-    q(x) = \begin{cases}
-      -1 & x < 0 \\\
-      1 & x \geq 0
-    \end{cases}
-    \\]
-
-    The gradient is estimated using the ApproxSign method.
-    \\[\frac{\partial q(x)}{\partial x} = \begin{cases}
-      (2 - 2 \left|x\right|) & \left|x\right| \leq 1 \\\
-      0 & \left|x\right| > 1
-    \end{cases}
-    \\]
-
-    ```plot-activation
-    quantizers.approx_sign
-    ```
-
-    # Arguments
-    x: Input tensor.
-
-    # Returns
-    Binarized tensor.
-
-    # References
-    - [Bi-Real Net: Enhancing the Performance of 1-bit CNNs With Improved
-      Representational Capability and Advanced
-      Training Algorithm](http://arxiv.org/abs/1808.00278)
-    """
-
     def grad(dy):
         abs_x = tf.math.abs(x)
         zeros = tf.zeros_like(dy)
@@ -238,38 +97,7 @@ def grad(dy):
     return math.sign(x), grad
 
 
-@utils.register_keras_custom_object
-@utils.set_precision(1)
 def swish_sign(x: tf.Tensor, beta: float = 5.0) -> tf.Tensor:
-    r"""Sign binarization function.
-
-    \\[
-    q(x) = \begin{cases}
-      -1 & x < 0 \\\
-      1 & x \geq 0
-    \end{cases}
-    \\]
-
-    The gradient is estimated using the SignSwish method.
-
-    \\[
-    \frac{\partial q_{\beta}(x)}{\partial x} = \frac{\beta\left\\{2-\beta x \tanh \left(\frac{\beta x}{2}\right)\right\\}}{1+\cosh (\beta x)}
-    \\]
-
-    ```plot-activation
-    quantizers.swish_sign
-    ```
-    # Arguments
-    x: Input tensor.
-    beta: Larger values result in a closer approximation to the derivative of the sign.
-
-    # Returns
-    Binarized tensor.
-
-    # References
-    - [BNN+: Improved Binary Network Training](https://arxiv.org/abs/1812.11800)
-    """
-
     @tf.custom_gradient
     def _call(x):
         def grad(dy):
@@ -281,58 +109,12 @@ def grad(dy):
     return _call(x)
 
 
-@utils.register_keras_custom_object
-@utils.set_precision(2)
 def ste_tern(
     x: tf.Tensor,
     threshold_value: float = 0.05,
     ternary_weight_networks: bool = False,
     clip_value: float = 1.0,
 ) -> tf.Tensor:
-    r"""Ternarization function.
-
-    \\[
-    q(x) = \begin{cases}
-    +1 & x > \Delta \\\
-    0 & |x| < \Delta \\\
-     -1 & x < - \Delta
-    \end{cases}
-    \\]
-
-    where \\(\Delta\\) is defined as the threshold and can be passed as an argument,
-    or can be calculated as per the Ternary Weight Networks original paper, such that
-
-    \\[
-    \Delta = \frac{0.7}{n} \sum_{i=1}^{n} |W_i|
-    \\]
-    where we assume that \\(W_i\\) is generated from a normal distribution.
-
-    The gradient is estimated using the Straight-Through Estimator
-    (essentially the Ternarization is replaced by a clipped identity on the
-    backward pass).
-    \\[\frac{\partial q(x)}{\partial x} = \begin{cases}
-    1 & \left|x\right| \leq \texttt{clip_value} \\\
-    0 & \left|x\right| > \texttt{clip_value}
-    \end{cases}\\]
-
-    ```plot-activation
-    quantizers.ste_tern
-    ```
-
-    # Arguments
-    x: Input tensor.
-    threshold_value: The value for the threshold, \\(\Delta\\).
-    ternary_weight_networks: Boolean of whether to use the
-        Ternary Weight Networks threshold calculation.
-    clip_value: Threshold for clipping gradients. If `None` gradients are not clipped.
-
-    # Returns
-    Ternarized tensor.
-
-    # References
-    - [Ternary Weight Networks](http://arxiv.org/abs/1605.04711)
-    """
-
     @tf.custom_gradient
     def _call(x):
         if ternary_weight_networks:
@@ -348,40 +130,7 @@ def grad(dy):
     return _call(x)
 
 
-@utils.register_keras_custom_object
-@utils.set_precision(1)
 def ste_heaviside(x: tf.Tensor, clip_value: float = 1.0) -> tf.Tensor:
-    r"""
-    Binarization function with output values 0 and 1.
-
-    \\[
-    q(x) = \begin{cases}
-    +1 & x > 0 \\\
-    0 & x \leq 0
-    \end{cases}
-    \\]
-
-    The gradient is estimated using the Straight-Through Estimator
-    (essentially the binarization is replaced by a clipped identity on the
-    backward pass).
-
-    \\[\frac{\partial q(x)}{\partial x} = \begin{cases}
-    1 & \left|x\right| \leq 1 \\\
-    0 & \left|x\right| > 1
-    \end{cases}\\]
-
-    ```plot-activation
-    quantizers.ste_heaviside
-    ```
-
-    # Arguments
-    x: Input tensor.
-    clip_value: Threshold for clipping gradients. If `None` gradients are not clipped.
-
-    # Returns
-    AND-binarized tensor.
-    """
-
     @tf.custom_gradient
     def _call(x):
         def grad(dy):
@@ -392,56 +141,54 @@ def grad(dy):
     return _call(x)
 
 
+@utils.register_alias("ste_sign")
 @utils.register_keras_custom_object
-@utils.set_precision(2)
-def dorefa_quantizer(x: tf.Tensor, k_bit: int = 2) -> tf.Tensor:
-    r"""k_bit quantizer as in the DoReFa paper.
+class SteSign(tf.keras.layers.Layer):
+    r"""Instantiates a serializable binary quantizer.
 
     \\[
     q(x) = \begin{cases}
-    0 & x < \frac{1}{2n} \\\
-    \frac{i}{n} & \frac{2i-1}{2n} < x < \frac{2i+1}{2n} \text{ for } i \in \\{1,n-1\\}\\\
-     1 & \frac{2n-1}{2n} < x
+      -1 & x < 0 \\\
+      1 & x \geq 0
     \end{cases}
     \\]
 
-    where \\(n = 2^{\text{k_bit}} - 1\\). The number of bits, k_bit, needs to be passed as an argument.
     The gradient is estimated using the Straight-Through Estimator
     (essentially the binarization is replaced by a clipped identity on the
     backward pass).
     \\[\frac{\partial q(x)}{\partial x} = \begin{cases}
-    1 &  0 \leq x \leq 1 \\\
-    0 & \text{else}
+      1 & \left|x\right| \leq \texttt{clip_value} \\\
+      0 & \left|x\right| > \texttt{clip_value}
     \end{cases}\\]
 
     ```plot-activation
-    quantizers.dorefa_quantizer
+    quantizers.SteSign
     ```
 
     # Arguments
-    k_bit: number of bits for the quantization.
-
-    # Returns
-    quantized tensor
+    clip_value: Threshold for clipping gradients. If `None` gradients are not clipped.
 
     # References
-    - [DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks
-      with Low Bitwidth Gradients](https://arxiv.org/abs/1606.06160)
+    - [Binarized Neural Networks: Training Deep Neural Networks with Weights and
+      Activations Constrained to +1 or -1](http://arxiv.org/abs/1602.02830)
     """
-    x = tf.clip_by_value(x, 0.0, 1.0)
+    precision = 1
 
-    @tf.custom_gradient
-    def _k_bit_with_identity_grad(x):
-        n = 2 ** k_bit - 1
-        return tf.round(x * n) / n, lambda dy: dy
+    def __init__(self, clip_value: float = 1.0, **kwargs):
+        self.clip_value = clip_value
+        super().__init__(**kwargs)
+
+    def call(self, inputs):
+        return ste_sign(inputs, clip_value=self.clip_value)
 
-    return _k_bit_with_identity_grad(x)
+    def get_config(self):
+        return {**super().get_config(), "clip_value": self.clip_value}
 
 
+@utils.register_alias("approx_sign")
 @utils.register_keras_custom_object
-class SteSign(QuantizerFunctionWrapper):
+class ApproxSign(tf.keras.layers.Layer):
     r"""Instantiates a serializable binary quantizer.
-
     \\[
     q(x) = \begin{cases}
       -1 & x < 0 \\\
@@ -449,32 +196,31 @@ class SteSign(QuantizerFunctionWrapper):
     \end{cases}
     \\]
 
-    The gradient is estimated using the Straight-Through Estimator
-    (essentially the binarization is replaced by a clipped identity on the
-    backward pass).
+    The gradient is estimated using the ApproxSign method.
     \\[\frac{\partial q(x)}{\partial x} = \begin{cases}
-      1 & \left|x\right| \leq \texttt{clip_value} \\\
-      0 & \left|x\right| > \texttt{clip_value}
-    \end{cases}\\]
+      (2 - 2 \left|x\right|) & \left|x\right| \leq 1 \\\
+      0 & \left|x\right| > 1
+    \end{cases}
+    \\]
 
     ```plot-activation
-    quantizers.ste_sign
+    quantizers.ApproxSign
     ```
 
-    # Arguments
-    clip_value: Threshold for clipping gradients. If `None` gradients are not clipped.
-
     # References
-    - [Binarized Neural Networks: Training Deep Neural Networks with Weights and
-      Activations Constrained to +1 or -1](http://arxiv.org/abs/1602.02830)
+    - [Bi-Real Net: Enhancing the Performance of 1-bit CNNs With Improved
+      Representational Capability and Advanced
+      Training Algorithm](http://arxiv.org/abs/1808.00278)
     """
+    precision = 1
 
-    def __init__(self, clip_value: float = 1.0):
-        super().__init__(ste_sign, clip_value=clip_value)
+    def call(self, inputs):
+        return approx_sign(inputs)
 
 
+@utils.register_alias("ste_heaviside")
 @utils.register_keras_custom_object
-class SteHeaviside(QuantizerFunctionWrapper):
+class SteHeaviside(tf.keras.layers.Layer):
     r"""
     Instantiates a binarization quantizer with output values 0 and 1.
     \\[
@@ -494,7 +240,7 @@ class SteHeaviside(QuantizerFunctionWrapper):
     \end{cases}\\]
 
     ```plot-activation
-    quantizers.ste_heaviside
+    quantizers.SteHeaviside
     ```
 
     # Arguments
@@ -503,13 +249,22 @@ class SteHeaviside(QuantizerFunctionWrapper):
     # Returns
     AND Binarization function
     """
+    precision = 1
+
+    def __init__(self, clip_value: float = 1.0, **kwargs):
+        self.clip_value = clip_value
+        super().__init__(**kwargs)
 
-    def __init__(self, clip_value: float = 1.0):
-        super().__init__(ste_heaviside, clip_value=clip_value)
+    def call(self, inputs):
+        return ste_heaviside(inputs, clip_value=self.clip_value)
+
+    def get_config(self):
+        return {**super().get_config(), "clip_value": self.clip_value}
 
 
+@utils.register_alias("swish_sign")
 @utils.register_keras_custom_object
-class SwishSign(QuantizerFunctionWrapper):
+class SwishSign(tf.keras.layers.Layer):
     r"""Sign binarization function.
 
     \\[
@@ -526,7 +281,7 @@ class SwishSign(QuantizerFunctionWrapper):
     \\]
 
     ```plot-activation
-    quantizers.swish_sign
+    quantizers.SwishSign
     ```
     # Arguments
     beta: Larger values result in a closer approximation to the derivative of the sign.
@@ -537,13 +292,22 @@ class SwishSign(QuantizerFunctionWrapper):
     # References
     - [BNN+: Improved Binary Network Training](https://arxiv.org/abs/1812.11800)
     """
+    precision = 1
+
+    def __init__(self, beta: float = 5.0, **kwargs):
+        self.beta = beta
+        super().__init__(**kwargs)
+
+    def call(self, inputs):
+        return swish_sign(inputs, beta=self.beta)
 
-    def __init__(self, beta: float = 5.0):
-        super().__init__(swish_sign, beta=beta)
+    def get_config(self):
+        return {**super().get_config(), "beta": self.beta}
 
 
+@utils.register_alias("magnitude_aware_sign")
 @utils.register_keras_custom_object
-class MagnitudeAwareSign(QuantizerFunctionWrapper):
+class MagnitudeAwareSign(tf.keras.layers.Layer):
     r"""Instantiates a serializable magnitude-aware sign quantizer for Bi-Real Net.
 
     A scaled sign function computed according to Section 3.3 in
@@ -562,13 +326,26 @@ class MagnitudeAwareSign(QuantizerFunctionWrapper):
       Algorithm](https://arxiv.org/abs/1808.00278)
 
     """
+    precision = 1
+
+    def __init__(self, clip_value: float = 1.0, **kwargs):
+        self.clip_value = clip_value
+        super().__init__(**kwargs)
+
+    def call(self, inputs):
+        scale_factor = tf.stop_gradient(
+            tf.reduce_mean(tf.abs(inputs), axis=list(range(len(inputs.shape) - 1)))
+        )
 
-    def __init__(self, clip_value: float = 1.0):
-        super().__init__(magnitude_aware_sign, clip_value=clip_value)
+        return scale_factor * ste_sign(inputs, clip_value=self.clip_value)
+
+    def get_config(self):
+        return {**super().get_config(), "clip_value": self.clip_value}
 
 
+@utils.register_alias("ste_tern")
 @utils.register_keras_custom_object
-class SteTern(QuantizerFunctionWrapper):
+class SteTern(tf.keras.layers.Layer):
     r"""Instantiates a serializable ternarization quantizer.
 
     \\[
@@ -596,7 +373,7 @@ class SteTern(QuantizerFunctionWrapper):
     \end{cases}\\]
 
     ```plot-activation
-    quantizers.ste_tern
+    quantizers.SteTern
     ```
 
     # Arguments
@@ -609,22 +386,40 @@ class SteTern(QuantizerFunctionWrapper):
     - [Ternary Weight Networks](http://arxiv.org/abs/1605.04711)
     """
 
+    precision = 2
+
     def __init__(
         self,
         threshold_value: float = 0.05,
         ternary_weight_networks: bool = False,
         clip_value: float = 1.0,
+        **kwargs,
     ):
-        super().__init__(
-            ste_tern,
-            threshold_value=threshold_value,
-            ternary_weight_networks=ternary_weight_networks,
-            clip_value=clip_value,
+        self.threshold_value = threshold_value
+        self.ternary_weight_networks = ternary_weight_networks
+        self.clip_value = clip_value
+        super().__init__(**kwargs)
+
+    def call(self, inputs):
+        return ste_tern(
+            inputs,
+            threshold_value=self.threshold_value,
+            ternary_weight_networks=self.ternary_weight_networks,
+            clip_value=self.clip_value,
         )
 
+    def get_config(self):
+        return {
+            **super().get_config(),
+            "threshold_value": self.threshold_value,
+            "ternary_weight_networks": self.ternary_weight_networks,
+            "clip_value": self.clip_value,
+        }
+
 
+@utils.register_alias("dorefa_quantizer")
 @utils.register_keras_custom_object
-class DoReFaQuantizer(QuantizerFunctionWrapper):
+class DoReFaQuantizer(tf.keras.layers.Layer):
     r"""Instantiates a serializable k_bit quantizer as in the DoReFa paper.
 
     \\[
@@ -645,7 +440,7 @@ class DoReFaQuantizer(QuantizerFunctionWrapper):
     \end{cases}\\]
 
     ```plot-activation
-    quantizers.dorefa_quantizer
+    quantizers.DoReFaQuantizer
     ```
 
     # Arguments
@@ -658,16 +453,30 @@ class DoReFaQuantizer(QuantizerFunctionWrapper):
     - [DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks
       with Low Bitwidth Gradients](https://arxiv.org/abs/1606.06160)
     """
+    precision = None
 
-    def __init__(self, k_bit: int):
-        super().__init__(dorefa_quantizer, k_bit=k_bit)
+    def __init__(self, k_bit: int = 2, **kwargs):
         self.precision = k_bit
+        super().__init__(**kwargs)
+
+    def call(self, inputs):
+        inputs = tf.clip_by_value(inputs, 0.0, 1.0)
+
+        @tf.custom_gradient
+        def _k_bit_with_identity_grad(x):
+            n = 2 ** self.precision - 1
+            return tf.round(x * n) / n, lambda dy: dy
+
+        return _k_bit_with_identity_grad(inputs)
+
+    def get_config(self):
+        return {**super().get_config(), "k_bit": self.precision}
 
 
-Quantizer = Union[QuantizerFunctionWrapper, Callable[[tf.Tensor], tf.Tensor]]
+Quantizer = Union[tf.keras.layers.Layer, Callable[[tf.Tensor], tf.Tensor]]
 
 
-def serialize(quantizer: Quantizer):
+def serialize(quantizer: tf.keras.layers.Layer):
     return tf.keras.utils.serialize_keras_object(quantizer)
 
 
diff --git a/plot_altair.py b/plot_altair.py
--- a/plot_altair.py
+++ b/plot_altair.py
@@ -1,3 +1,4 @@
+import inspect
 import os
 import uuid
 from functools import reduce
@@ -44,6 +45,8 @@ def html_format(source, language=None, css_class=None, options=None, md=None):
 
 def plot_activation(source, language=None, css_class=None, options=None, md=None):
     function = reduce(getattr, [lq, *source.split(".")])
+    if inspect.isclass(function):
+        function = function()
     x = np.linspace(-2, 2, 500)
     y, dy = calculate_activation(function, x)
     data = pd.DataFrame({"x": x, "y": y, "dy / dx": dy})


Test patch:
diff --git a/larq/quantized_variable_test.py b/larq/quantized_variable_test.py
--- a/larq/quantized_variable_test.py
+++ b/larq/quantized_variable_test.py
@@ -21,7 +21,7 @@ def test_inheritance(distribute_scope):
     assert isinstance(quantized_variable, DistributedVariable) is distribute_scope  # type: ignore
 
 
-def test_read(distribute_scope, eager_and_graph_mode):
+def test_read(eager_and_graph_mode, distribute_scope):
     x = QuantizedVariable.from_variable(get_var(3.5), quantizer=lambda x: 2 * x)
     evaluate(x.initializer)
 
@@ -48,7 +48,7 @@ def test_sparse_reads(eager_and_graph_mode):
         assert evaluate(x.gather_nd([0])) == 2
 
 
-def test_read_nested_scopes(distribute_scope, eager_and_graph_mode):
+def test_read_nested_scopes(eager_and_graph_mode, distribute_scope):
     x = QuantizedVariable.from_variable(get_var(3.5), quantizer=lambda x: 2 * x)
     evaluate(x.initializer)
     with quantized_scope.scope(True):
@@ -58,7 +58,7 @@ def test_read_nested_scopes(distribute_scope, eager_and_graph_mode):
         assert evaluate(x.read_value()) == 7
 
 
-def test_method_delegations(distribute_scope, eager_and_graph_mode):
+def test_method_delegations(eager_and_graph_mode, distribute_scope):
     x = QuantizedVariable.from_variable(get_var(3.5), quantizer=lambda x: 2 * x)
     with quantized_scope.scope(True):
         evaluate(x.initializer)
@@ -118,7 +118,7 @@ def slices(val, index):
         )
 
 
-def test_overloads(quantized, distribute_scope, eager_and_graph_mode):
+def test_overloads(eager_and_graph_mode, quantized, distribute_scope):
     if quantized:
         x = QuantizedVariable.from_variable(get_var(3.5), quantizer=lambda x: 2 * x)
     else:
@@ -167,7 +167,7 @@ def test_tensor_equality(quantized, eager_mode):
         assert_array_equal(x != [7.0, 8.0, 10.0], [False, False, True])
 
 
-def test_assign(quantized, distribute_scope, eager_and_graph_mode):
+def test_assign(eager_and_graph_mode, quantized, distribute_scope):
     x = QuantizedVariable.from_variable(
         get_var(0.0, tf.float64), quantizer=lambda x: 2 * x
     )
diff --git a/larq/quantizers_test.py b/larq/quantizers_test.py
--- a/larq/quantizers_test.py
+++ b/larq/quantizers_test.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 import tensorflow as tf
+from packaging import version
 
 import larq as lq
 from larq import testing_utils
@@ -22,37 +23,32 @@ class TestCommonFunctionality:
 
     @pytest.mark.parametrize("module", [lq.quantizers, tf.keras.activations])
     @pytest.mark.parametrize(
-        "name",
-        ["ste_sign", "approx_sign", "magnitude_aware_sign", "swish_sign", "ste_tern"],
+        "name,ref_cls",
+        [
+            ("ste_sign", lq.quantizers.SteSign),
+            ("approx_sign", lq.quantizers.ApproxSign),
+            ("ste_heaviside", lq.quantizers.SteHeaviside),
+            ("magnitude_aware_sign", lq.quantizers.MagnitudeAwareSign),
+            ("swish_sign", lq.quantizers.SwishSign),
+            ("ste_tern", lq.quantizers.SteTern),
+        ],
     )
-    def test_serialization(self, module, name):
+    def test_serialization(self, module, name, ref_cls):
         fn = module.get(name)
-        ref_fn = getattr(lq.quantizers, name)
-        assert fn == ref_fn
+        assert fn.__class__ == ref_cls
+        fn = module.get(ref_cls())
+        assert fn.__class__ == ref_cls
         assert type(fn.precision) == int
+        if module == tf.keras.activations and version.parse(
+            tf.__version__
+        ) < version.parse("1.15"):
+            pytest.skip(
+                "TensorFlow < 1.15 does not support Quantizer classes as activations"
+            )
         config = module.serialize(fn)
         fn = module.deserialize(config)
-        assert fn == ref_fn
+        assert fn.__class__ == ref_cls
         assert type(fn.precision) == int
-        fn = module.get(ref_fn)
-        assert fn == ref_fn
-        assert type(fn.precision) == int
-
-    @pytest.mark.parametrize(
-        "ref_fn",
-        [
-            lq.quantizers.SteSign(),
-            lq.quantizers.SteHeaviside(),
-            lq.quantizers.MagnitudeAwareSign(),
-            lq.quantizers.SwishSign(),
-            lq.quantizers.SteTern(),
-        ],
-    )
-    def test_serialization_cls(self, ref_fn):
-        assert type(ref_fn.precision) == int
-        config = lq.quantizers.serialize(ref_fn)
-        fn = lq.quantizers.deserialize(config)
-        assert fn.__class__ == ref_fn.__class__
 
     def test_invalid_usage(self):
         with pytest.raises(ValueError):
@@ -84,6 +80,7 @@ class TestQuantization:
             "ste_sign",
             lq.quantizers.SteSign(),
             "approx_sign",
+            lq.quantizers.ApproxSign(),
             "swish_sign",
             lq.quantizers.SwishSign(),
         ],
@@ -122,7 +119,7 @@ def test_and_binarization(self, fn):
     def test_magnitude_aware_sign_binarization(self, eager_mode):
         a = np.random.uniform(-2, 2, (3, 2, 2, 3))
         x = tf.Variable(a)
-        y = lq.quantizers.magnitude_aware_sign(x)
+        y = lq.quantizers.MagnitudeAwareSign()(x)
 
         assert y.shape == x.shape
 
@@ -207,12 +204,9 @@ def test_ternarization_with_ternary_weight_networks(self):
         assert not np.any(result > 1)
         assert not np.any(result < -1)
 
-    @pytest.mark.parametrize(
-        "fn", [lq.quantizers.dorefa_quantizer, lq.quantizers.DoReFaQuantizer(2)]
-    )
-    def test_dorefa_quantize(self, fn):
+    def test_dorefa_quantize(self):
         x = tf.keras.backend.placeholder(ndim=2)
-        f = tf.keras.backend.function([x], [fn(x)])
+        f = tf.keras.backend.function([x], [lq.quantizers.DoReFaQuantizer(2)(x)])
         real_values = testing_utils.generate_real_values_with_zeros()
         result = f([real_values])[0]
         k_bit = 2
@@ -234,19 +228,27 @@ class TestGradients:
 
     @pytest.mark.parametrize(
         "fn",
-        [lq.quantizers.ste_sign, lq.quantizers.ste_tern, lq.quantizers.ste_heaviside],
+        [
+            lq.quantizers.SteSign(clip_value=None),
+            lq.quantizers.SteTern(clip_value=None),
+            lq.quantizers.SteHeaviside(clip_value=None),
+        ],
     )
     def test_identity_ste_grad(self, eager_mode, fn):
         x = testing_utils.generate_real_values_with_zeros(shape=(8, 3, 3, 16))
         tf_x = tf.Variable(x)
         with tf.GradientTape() as tape:
-            activation = fn(tf_x, clip_value=None)
+            activation = fn(tf_x)
         grad = tape.gradient(activation, tf_x)
         np.testing.assert_allclose(grad.numpy(), np.ones_like(x))
 
     @pytest.mark.parametrize(
         "fn",
-        [lq.quantizers.ste_sign, lq.quantizers.ste_tern, lq.quantizers.ste_heaviside],
+        [
+            lq.quantizers.SteSign(),
+            lq.quantizers.SteTern(),
+            lq.quantizers.SteHeaviside(),
+        ],
     )
     def test_ste_grad(self, eager_mode, fn):
         @np.vectorize
@@ -272,12 +274,12 @@ def swish_grad(x, beta):
         x = testing_utils.generate_real_values_with_zeros(shape=(8, 3, 3, 16))
         tf_x = tf.Variable(x)
         with tf.GradientTape() as tape:
-            activation = lq.quantizers.swish_sign(tf_x)
+            activation = lq.quantizers.SwishSign()(tf_x)
         grad = tape.gradient(activation, tf_x)
         np.testing.assert_allclose(grad.numpy(), swish_grad(x, beta=5.0))
 
         with tf.GradientTape() as tape:
-            activation = lq.quantizers.swish_sign(tf_x, beta=10.0)
+            activation = lq.quantizers.SwishSign(beta=10.0)(tf_x)
         grad = tape.gradient(activation, tf_x)
         np.testing.assert_allclose(grad.numpy(), swish_grad(x, beta=10.0))
 
@@ -291,7 +293,7 @@ def approx_sign_grad(x):
         x = testing_utils.generate_real_values_with_zeros(shape=(8, 3, 3, 16))
         tf_x = tf.Variable(x)
         with tf.GradientTape() as tape:
-            activation = lq.quantizers.approx_sign(tf_x)
+            activation = lq.quantizers.ApproxSign()(tf_x)
         grad = tape.gradient(activation, tf_x)
         np.testing.assert_allclose(grad.numpy(), approx_sign_grad(x))
 
@@ -299,7 +301,7 @@ def test_magnitude_aware_sign_grad(self, eager_mode):
         a = np.random.uniform(-2, 2, (3, 2, 2, 3))
         x = tf.Variable(a)
         with tf.GradientTape() as tape:
-            y = lq.quantizers.magnitude_aware_sign(x)
+            y = lq.quantizers.MagnitudeAwareSign()(x)
         grad = tape.gradient(y, x)
 
         scale_vector = [