"""A Quantizer defines the way of transforming a full precision input to a
quantized output and the pseudo-gradient method used for the backwards pass."""

import tensorflow as tf
from larq import utils


def sign(x):
    """A sign function that will never be zero"""
    return tf.sign(tf.sign(x) + 0.1)


@tf.custom_gradient
def _binarize_with_identity_grad(x):
    def grad(dy):
        return dy

    return sign(x), grad


@tf.custom_gradient
def _binarize_with_weighted_grad(x):
    def grad(dy):
        return (1 - tf.abs(x)) * 2 * dy

    return sign(x), grad


@utils.register_keras_custom_object
def ste_sign(x):
    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 1 \\\
      0 & \left|x\right| > 1
    \end{cases}\\]

    # Arguments
    x: Input tensor.

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

    x = tf.clip_by_value(x, -1, 1)

    return _binarize_with_identity_grad(x)


@utils.register_keras_custom_object
def magnitude_aware_sign(x):
    r"""
    Magnitude-aware sign for birealnet.


    # Arguments
    x: Input tensor

    # 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.stop_gradient(
        tf.reduce_mean(tf.abs(x), axis=list(range(len(x.shape) - 1)))
    )
    return scale_factor * ste_sign(x)


@utils.register_keras_custom_object
def approx_sign(x):
    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}
    \\]

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

    x = tf.clip_by_value(x, -1, 1)

    return _binarize_with_weighted_grad(x)


def serialize(initializer):
    return tf.keras.utils.serialize_keras_object(initializer)


def deserialize(name, custom_objects=None):
    return tf.keras.utils.deserialize_keras_object(
        name,
        module_objects=globals(),
        custom_objects=custom_objects,
        printable_module_name="quantization function",
    )


def get(identifier):
    if identifier is None:
        return None
    if isinstance(identifier, str):
        return deserialize(str(identifier))
    if callable(identifier):
        return identifier
    raise ValueError(
        f"Could not interpret quantization function identifier: {identifier}"
    )

larq/quantizers.py
"""A Quantizer defines the way of transforming a full precision input to a
quantized output and the pseudo-gradient method used for the backwards pass."""

import tensorflow as tf
from larq import utils


def sign(x):
    """A sign function that will never be zero"""
    return tf.sign(tf.sign(x) + 0.1)


@tf.custom_gradient
def _binarize_with_identity_grad(x):
    def grad(dy):
        return dy

    return sign(x), grad


@tf.custom_gradient
def _binarize_with_weighted_grad(x):
    def grad(dy):
        return (1 - tf.abs(x)) * 2 * dy

    return sign(x), grad


@utils.register_keras_custom_object
def ste_sign(x):
    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 1 \\\
      0 & \left|x\right| > 1
    \end{cases}\\]

    # Arguments
    x: Input tensor.

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

    x = tf.clip_by_value(x, -1, 1)

    return _binarize_with_identity_grad(x)


@utils.register_keras_custom_object
def magnitude_aware_sign(x):
    r"""
    Magnitude-aware sign for Bi-Real Net.


    # Arguments
    x: Input tensor

    # 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.stop_gradient(
        tf.reduce_mean(tf.abs(x), axis=list(range(len(x.shape) - 1)))
    )
    return scale_factor * ste_sign(x)


@utils.register_keras_custom_object
def approx_sign(x):
    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}
    \\]

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

    x = tf.clip_by_value(x, -1, 1)

    return _binarize_with_weighted_grad(x)


def serialize(initializer):
    return tf.keras.utils.serialize_keras_object(initializer)


def deserialize(name, custom_objects=None):
    return tf.keras.utils.deserialize_keras_object(
        name,
        module_objects=globals(),
        custom_objects=custom_objects,
        printable_module_name="quantization function",
    )


def get(identifier):
    if identifier is None:
        return None
    if isinstance(identifier, str):
        return deserialize(str(identifier))
    if callable(identifier):
        return identifier
    raise ValueError(
        f"Could not interpret quantization function identifier: {identifier}"
    )

larq/quantizers.py
diff --git a/README.md b/README.md
index fd8f466d..16450c41 100644
--- a/README.md
+++ b/README.md
@@ -10,17 +10,21 @@ Existing deep neural networks use 32 bits, 16 bits or 8 bits to encode each weig
 
 ## Getting Started
 
-To build a QNN, Larq introduces the concept of _Quantized Layers_ and _Quantizers_. A _Quantizer_ defines the way of transforming a full precision input to a quantized output and the pseudo-gradient method used for the backwards pass. Each _Quantized Layer_ requires a `kernel_quantizer` and an `input_quantizer` that describe the way of quantizing the weights of the layer and the activations of the previous layer respectively. If both `input_quantizer` and `kernel_quantizer` are `None` the layer is equivalent to a full precision layer.
+To build a QNN, Larq introduces the concept of [quantized layers](https://plumerai.github.io/larq/api/layers/) and [quantizers](https://plumerai.github.io/larq/api/quantizers/). A quantizer defines the way of transforming a full precision input to a quantized output and the pseudo-gradient method used for the backwards pass. Each quantized layer requires a `kernel_quantizer` and an `input_quantizer` that describe the way of quantizing the weights of the layer and the activations of the previous layer respectively. If both `input_quantizer` and `kernel_quantizer` are `None` the layer is equivalent to a full precision layer.
 
-You can define a binarized densely-connected layer using the Straight-Through Estimator the following way:
+You can define a simple binarized fully-connected Keras model using the [Straight-Through Estimator](https://plumerai.github.io/larq/api/quantizers/#ste_sign) the following way:
 
 ```python
-larq.layers.QuantDense(
-    512,
-    input_quantizer="ste_sign",
-    kernel_quantizer="ste_sign",
-    kernel_constraint="weight_clip",
-)
+model = tf.keras.models.Sequential([
+    tf.keras.layers.Flatten(),
+    larq.layers.QuantDense(512,
+                           kernel_quantizer="ste_sign",
+                           kernel_constraint="weight_clip"),
+    larq.layers.QuantDense(10,
+                           input_quantizer="ste_sign",
+                           kernel_quantizer="ste_sign",
+                           kernel_constraint="weight_clip",
+                           activation="softmax")])
 ```
 
 This layer can be used inside a [Keras model](https://www.tensorflow.org/alpha/guide/keras/overview#sequential_model) or with a [custom training loop](https://www.tensorflow.org/alpha/guide/keras/overview#model_subclassing).
diff --git a/apidocs.yml b/apidocs.yml
index 95220ced..5f67c561 100644
--- a/apidocs.yml
+++ b/apidocs.yml
@@ -24,3 +24,4 @@
     - larq.quantizers:
         - larq.quantizers.ste_sign
         - larq.quantizers.approx_sign
+        - larq.quantizers.magnitude_aware_sign
diff --git a/docs/examples/binarynet_advanced_cifar10.ipynb b/docs/examples/binarynet_advanced_cifar10.ipynb
index 2c8248e3..8599082e 100644
--- a/docs/examples/binarynet_advanced_cifar10.ipynb
+++ b/docs/examples/binarynet_advanced_cifar10.ipynb
@@ -96,7 +96,7 @@
     {
       "cell_type": "markdown",
       "source": [
-        "## Build Binarynet\n",
+        "## Build BinaryNet\n",
         "\n",
         "Here we build the binarynet model layer by layer using a keras sequential model:"
       ],
diff --git a/docs/guide.md b/docs/guide.md
new file mode 100644
index 00000000..294645e5
--- /dev/null
+++ b/docs/guide.md
@@ -0,0 +1,94 @@
+# User Guide
+
+To build a Quantized Neural Network (QNN), Larq introduces the concept of [quantized layers](https://plumerai.github.io/larq/api/layers/) and [quantizers](https://plumerai.github.io/larq/api/quantizers/). A quantizer defines the way of transforming a full precision input to a quantized output and the pseudo-gradient method used for the backwards pass.
+
+Each quantized layer requires a `kernel_quantizer` and an `input_quantizer` that describe the way of quantizing the weights of the layer and the activations of the previous layer respectively. If both `input_quantizer` and `kernel_quantizer` are `None` the layer is equivalent to a full precision layer. Larq layers are fully compatible with the Keras API so you can use them with Keras Layers interchangeably:
+
+```python tab="Larq 32-bit model"
+model = tf.keras.models.Sequential([
+    tf.keras.layers.Flatten(),
+    larq.layers.QuantDense(512, activation="relu"),
+    larq.layers.QuantDense(10, activation="softmax")
+])
+```
+
+```python tab="Keras 32-bit model"
+model = tf.keras.models.Sequential([
+    tf.keras.layers.Flatten(),
+    tf.keras.layers.Dense(512, activation="relu"),
+    tf.keras.layers.Dense(10, activation="softmax")
+])
+```
+
+A simple fully-connected Binarized Neural Network (BNN) using the [Straight-Through Estimator](https://plumerai.github.io/larq/api/quantizers/#ste_sign) can be defined in just a few lines of code using either the Keras sequential, functional or model subclassing APIs:
+
+```python tab="Larq 1-bit model"
+model = tf.keras.models.Sequential([
+    tf.keras.layers.Flatten(),
+    larq.layers.QuantDense(512,
+                           kernel_quantizer="ste_sign",
+                           kernel_constraint="weight_clip"),
+    larq.layers.QuantDense(10,
+                           input_quantizer="ste_sign",
+                           kernel_quantizer="ste_sign",
+                           kernel_constraint="weight_clip",
+                           activation="softmax")])
+```
+
+```python tab="Larq 1-bit model functional"
+x = tf.keras.Input(shape=(28, 28, 1))
+y = tf.keras.layers.Flatten()(x)
+y = larq.layers.QuantDense(512,
+                           kernel_quantizer="ste_sign",
+                           kernel_constraint="weight_clip")(y)
+y = larq.layers.QuantDense(10,
+                           input_quantizer="ste_sign",
+                           kernel_quantizer="ste_sign",
+                           kernel_constraint="weight_clip",
+                           activation="softmax")(y)
+model = tf.keras.Model(inputs=x, outputs=y)
+```
+
+```python tab="Larq 1-bit model subclassing"
+class MyModel(tf.keras.Model):
+    def __init__(self):
+        super().__init__()
+        self.flatten = tf.keras.layers.Flatten()
+        self.dense1 = larq.layers.QuantDense(512,
+                                             kernel_quantizer="ste_sign",
+                                             kernel_constraint="weight_clip")
+        self.dense2 = larq.layers.QuantDense(10,
+                                             input_quantizer="ste_sign",
+                                             kernel_quantizer="ste_sign",
+                                             kernel_constraint="weight_clip",
+                                             activation="softmax")
+
+    def call(self, inputs):
+        x = self.flatten(inputs)
+        x = self.dense1(x)
+        return self.dense2(x)
+
+model = MyModel()
+```
+
+## Using Custom Quantizers
+
+Quantizers are functions that transform a full precision input to a quantized output. Since this transformation usually is non-differentiable it is necessary to modify the gradient in order to be able to train the resulting QNN. This can be done with the [`tf.custom_gradient`](https://www.tensorflow.org/api_docs/python/tf/custom_gradient) decorator.
+
+In this example we will define a binarization function with an identity gradient:
+
+```python
+@tf.custom_gradient
+def identity_sign(x):
+    def grad(dy):
+        return dy
+    return tf.sign(x), grad
+```
+
+This function can now be used as an `input_quantizer` or a `kernel_quantizer`:
+
+```python
+larq.layers.QuantDense(10,
+                       input_quantizer=identity_sign,
+                       kernel_quantizer=identity_sign)
+```
diff --git a/larq/quantizers.py b/larq/quantizers.py
index d2c984ad..bb1b62ac 100644
--- a/larq/quantizers.py
+++ b/larq/quantizers.py
@@ -64,7 +64,7 @@ def ste_sign(x):
 @utils.register_keras_custom_object
 def magnitude_aware_sign(x):
     r"""
-    Magnitude-aware sign for birealnet.
+    Magnitude-aware sign for Bi-Real Net.
 
 
     # Arguments
diff --git a/mkdocs.yml b/mkdocs.yml
index 2603af87..abf47471 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -4,14 +4,15 @@ site_description: "An Open Source Machine Learning Framework for Training Extrem
 nav:
   - Home:
       - Getting Started: index.md
-      - Library Development:
+      - User Guide: guide.md
+      - Developer Guide:
           - Contributing Guide: contributing.md
           - Code of Conduct: code_of_conduct.md
   - Examples:
       - Introduction to BNNs with Larq: examples/mnist.ipynb
-      - Binarynet on CIFAR10: examples/binarynet_cifar10.ipynb
-      - Binarynet on CIFAR10 (Advanced): examples/binarynet_advanced_cifar10.ipynb
-  - Documentation:
+      - BinaryNet on CIFAR10: examples/binarynet_cifar10.ipynb
+      - BinaryNet on CIFAR10 (Advanced): examples/binarynet_advanced_cifar10.ipynb
+  - API:
       - Layers: api/layers.md
       - Quantizers: api/quantizers.md
       - Activations: api/activations.md

Add docs on how to define your own quantizer

