Repository: tensorflow/addons

Issue:
Request for example: CRF
This is pending #22 and #314 but once it's merged we should be sure to add some examples as it's highly requested.



Discussion:
Harsh188: I have no idea what's wrong with this. When I run `pytest -v --doctest-modules tensorflow_addons/layers/crf.py` I get no errors.

Patch:
diff --git a/docs/README.md b/docs/README.md
index ce58a50..9ce374a 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -21,22 +21,11 @@ python -m tensorflow_docs.tools.nbfmt {path to notebook file or directory}
 
 
 
-## 2. Generated API docs
 
-[tensorflow.org/addons/api_docs/python/tfa](https://tensorflow.org/addons/api_docs/python/tfa)
 
-`build_docs.py` controls executed this docs generation. To test-run it:
 
-```bash
-# Install dependencies:
-pip install -r tools/install_deps/doc_requirements.txt
 
-# Build tool:
-bazel build docs:build_docs
 
-# Generate API doc:
-# Use current branch
-bazel-bin/docs/build_docs --git_branch=$(git rev-parse --abbrev-ref HEAD)
-# or specified explicitly
-bazel-bin/docs/build_docs --git_branch=master --output_dir=/tmp/tfa_api
-```
+
+
+
diff --git a/tensorflow_addons/layers/__init__.py b/tensorflow_addons/layers/__init__.py
index 7c89f95..78c972e 100644
--- a/tensorflow_addons/layers/__init__.py
+++ b/tensorflow_addons/layers/__init__.py
@@ -40,3 +40,4 @@ from tensorflow_addons.layers.wrappers import WeightNormalization
 from tensorflow_addons.layers.esn import ESN
 from tensorflow_addons.layers.stochastic_depth import StochasticDepth
 from tensorflow_addons.layers.noisy_dense import NoisyDense
+from tensorflow_addons.layers.crf import CRF
diff --git a/tensorflow_addons/layers/crf.py b/tensorflow_addons/layers/crf.py
index 22b3f9a..129a6bc 100644
--- a/tensorflow_addons/layers/crf.py
+++ b/tensorflow_addons/layers/crf.py
@@ -27,8 +27,38 @@ from tensorflow_addons.utils import types
 class CRF(tf.keras.layers.Layer):
     """Linear chain conditional random field (CRF).
 
+    Inherits from: `tf.keras.layers.Layer`.
+
     References:
         - [Conditional Random Field](https://en.wikipedia.org/wiki/Conditional_random_field)
+
+    Example:
+
+    >>> layer = tfa.layers.CRF(4)
+    >>> inputs = np.random.rand(2, 4, 8).astype(np.float32)
+    >>> decoded_sequence, potentials, sequence_length, chain_kernel = layer(inputs)
+    >>> decoded_sequence.shape
+    TensorShape([2, 4])
+    >>> potentials.shape
+    TensorShape([2, 4, 4])
+    >>> sequence_length
+    <tf.Tensor: shape=(2,), dtype=int64, numpy=array([4, 4])>
+    >>> chain_kernel.shape
+    TensorShape([4, 4])
+
+    Arguments:
+        units: Positive integer, dimensionality of the reservoir.
+        chain_initializer: Orthogonal matrix. Default to `orthogonal`.
+        use_boundary: `Boolean`, whether the layer uses a boundary vector. Default to `True`.
+        boundary_initializer: Tensors initialized to 0. Default to `zeros`.
+        use_kernel: `Boolean`, whether the layer uses a kernel weights. Default to `True`.
+    Call Arguments:
+        inputs: Positive integer, dimensionality of the output space.
+        mask: A boolean `Tensor` of shape `[batch_size, sequence_length]`
+            or `None`. Default to `None`.
+    Raises:
+        ValueError: If input mask doesn't have dim 2 or None.
+        NotImplementedError: If left padding is provided.
     """
 
     @typechecked
diff --git a/tensorflow_addons/metrics/cohens_kappa.py b/tensorflow_addons/metrics/cohens_kappa.py
index cb47353..26cc068 100644
--- a/tensorflow_addons/metrics/cohens_kappa.py
+++ b/tensorflow_addons/metrics/cohens_kappa.py
@@ -28,7 +28,7 @@ from typing import Optional
 class CohenKappa(Metric):
     """Computes Kappa score between two raters.
 
-    The score lies in the range [-1, 1]. A score of -1 represents
+    The score lies in the range `[-1, 1]`. A score of -1 represents
     complete disagreement between two raters whereas a score of 1
     represents complete agreement between the two raters.
     A score of 0 means agreement by chance.
@@ -36,30 +36,61 @@ class CohenKappa(Metric):
     Note: As of now, this implementation considers all labels
     while calculating the Cohen's Kappa score.
 
-    Usage:
-
-    ```python
-    actuals = np.array([4, 4, 3, 4, 2, 4, 1, 1], dtype=np.int32)
-    preds = np.array([4, 4, 3, 4, 4, 2, 1, 1], dtype=np.int32)
-    weights = np.array([1, 1, 2, 5, 10, 2, 3, 3], dtype=np.int32)
-
-    m = tfa.metrics.CohenKappa(num_classes=5, sparse_labels=True)
-    m.update_state(actuals, preds)
-    print('Final result: ', m.result().numpy()) # Result: 0.61904764
+    Args:
+        num_classes: Number of unique classes in your dataset.
+        weightage: (optional) Weighting to be considered for calculating
+            kappa statistics. A valid value is one of
+            [None, 'linear', 'quadratic']. Defaults to `None`
+        sparse_labels: (bool) Valid only for multi-class scenario.
+            If True, ground truth labels are expected tp be integers
+            and not one-hot encoded.
+        regression: (bool) If set, that means the problem is being treated
+            as a regression problem where you are regressing the predictions.
+            **Note:** If you are regressing for the values, the the output layer
+            should contain a single unit.
+        name: (optional) String name of the metric instance
+        dtype: (optional) Data type of the metric result. Defaults to `None`.
 
-    # To use this with weights, sample_weight argument can be used.
-    m = tfa.metrics.CohenKappa(num_classes=5, sparse_labels=True)
-    m.update_state(actuals, preds, sample_weight=weights)
-    print('Final result: ', m.result().numpy()) # Result: 0.37209308
-    ```
+    Raises:
+        ValueError: If the value passed for `weightage` is invalid
+        i.e. not any one of [None, 'linear', 'quadratic'].
 
-    Usage with tf.keras API:
+    Usage:
 
-    ```python
-    model = tf.keras.models.Model(inputs, outputs)
-    model.add_metric(tfa.metrics.CohenKappa(num_classes=5)(outputs))
-    model.compile('sgd', loss='mse')
-    ```
+    >>> y_true = np.array([4, 4, 3, 4, 2, 4, 1, 1], dtype=np.int32)
+    >>> y_pred = np.array([4, 4, 3, 4, 4, 2, 1, 1], dtype=np.int32)
+    >>> weights = np.array([1, 1, 2, 5, 10, 2, 3, 3], dtype=np.int32)
+    >>> metric = tfa.metrics.CohenKappa(num_classes=5, sparse_labels=True)
+    >>> metric.update_state(y_true , y_pred)
+    <tf.Tensor: shape=(5, 5), dtype=float32, numpy=
+     array([[0., 0., 0., 0., 0.],
+            [0., 2., 0., 0., 0.],
+            [0., 0., 0., 0., 1.],
+            [0., 0., 0., 1., 0.],
+            [0., 0., 1., 0., 3.]], dtype=float32)>
+    >>> result = metric.result()
+    >>> result.numpy()
+    0.61904764
+    >>> # To use this with weights, sample_weight argument can be used.
+    >>> metric = tfa.metrics.CohenKappa(num_classes=5, sparse_labels=True)
+    >>> metric.update_state(y_true , y_pred , sample_weight=weights)
+    <tf.Tensor: shape=(5, 5), dtype=float32, numpy=
+     array([[ 0.,  0.,  0.,  0.,  0.],
+            [ 0.,  6.,  0.,  0.,  0.],
+           [ 0.,  0.,  0.,  0., 10.],
+            [ 0.,  0.,  0.,  2.,  0.],
+            [ 0.,  0.,  2.,  0.,  7.]], dtype=float32)>
+    >>> result = metric.result()
+    >>> result.numpy()
+     0.37209308
+
+    Usage with `tf.keras` API:
+
+    >>> inputs = tf.keras.Input(shape=(10,))
+    >>> x = tf.keras.layers.Dense(10)(inputs)
+    >>> outputs = tf.keras.layers.Dense(1)(x)
+    >>> model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
+    >>> model.compile('sgd', loss='mse', metrics=[tfa.metrics.CohenKappa(num_classes=3, sparse_labels=True)])
     """
 
     @typechecked
@@ -72,27 +103,7 @@ class CohenKappa(Metric):
         regression: bool = False,
         dtype: AcceptableDTypes = None,
     ):
-        """Creates a `CohenKappa` instance.
-
-        Args:
-          num_classes: Number of unique classes in your dataset.
-          weightage: (optional) Weighting to be considered for calculating
-            kappa statistics. A valid value is one of
-            [None, 'linear', 'quadratic']. Defaults to `None`
-          sparse_labels: (bool) Valid only for multi-class scenario.
-            If True, ground truth labels are expected tp be integers
-            and not one-hot encoded
-          regression: (bool) If set, that means the problem is being treated
-            as a regression problem where you are regressing the predictions.
-            **Note:** If you are regressing for the values, the the output layer
-            should contain a single unit.
-          name: (optional) String name of the metric instance
-          dtype: (optional) Data type of the metric result. Defaults to `None`
-
-        Raises:
-          ValueError: If the value passed for `weightage` is invalid
-            i.e. not any one of [None, 'linear', 'quadratic']
-        """
+        """Creates a `CohenKappa` instance."""
         super().__init__(name=name, dtype=dtype)
 
         if weightage not in (None, "linear", "quadratic"):
diff --git a/tensorflow_addons/metrics/f_scores.py b/tensorflow_addons/metrics/f_scores.py
index 28c6b7d..4669e8b 100755
--- a/tensorflow_addons/metrics/f_scores.py
+++ b/tensorflow_addons/metrics/f_scores.py
@@ -24,13 +24,15 @@ from typing import Optional
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
 class FBetaScore(tf.keras.metrics.Metric):
-    """Computes F-Beta score.
+    r"""Computes F-Beta score.
 
     It is the weighted harmonic mean of precision
-    and recall. Output range is [0, 1]. Works for
+    and recall. Output range is `[0, 1]`. Works for
     both multi-class and multi-label classification.
 
-    F-Beta = (1 + beta^2) * (prec * recall) / ((beta^2 * prec) + recall)
+    $$
+    F_{\beta} = (1 + \beta^2) * \frac{\textrm{precision} * \textrm{precision}}{(\beta^2 \cdot \textrm{precision}) + \textrm{recall}}
+    $$
 
     Args:
         num_classes: Number of unique classes in the dataset.
@@ -43,13 +45,15 @@ class FBetaScore(tf.keras.metrics.Metric):
         threshold: Elements of `y_pred` greater than threshold are
             converted to be 1, and the rest 0. If threshold is
             None, the argmax is converted to 1, and the rest 0.
+        name: (Optional) String name of the metric instance.
+        dtype: (Optional) Data type of the metric result.
 
     Returns:
-        F-Beta Score: float
+        F-Beta Score: float.
 
     Raises:
         ValueError: If the `average` has values other than
-        [None, micro, macro, weighted].
+        `[None, 'micro', 'macro', 'weighted']`.
 
         ValueError: If the `beta` value is less than or equal
         to 0.
@@ -67,6 +71,20 @@ class FBetaScore(tf.keras.metrics.Metric):
         weighted: Metrics are computed for each class
             and returns the mean weighted by the
             number of true instances in each class.
+
+    Usage:
+
+    >>> metric = tfa.metrics.FBetaScore(num_classes=3, beta=2.0, threshold=0.5)
+    >>> y_true = np.array([[1, 1, 1],
+    ...                    [1, 0, 0],
+    ...                    [1, 1, 0]], np.int32)
+    >>> y_pred = np.array([[0.2, 0.6, 0.7],
+    ...                    [0.2, 0.6, 0.6],
+    ...                    [0.6, 0.8, 0.0]], np.float32)
+    >>> metric.update_state(y_true, y_pred)
+    >>> result = metric.result()
+    >>> result.numpy()
+    array([0.3846154 , 0.90909094, 0.8333334 ], dtype=float32)
     """
 
     @typechecked
@@ -85,7 +103,7 @@ class FBetaScore(tf.keras.metrics.Metric):
         if average not in (None, "micro", "macro", "weighted"):
             raise ValueError(
                 "Unknown average type. Acceptable values "
-                "are: [None, micro, macro, weighted]"
+                "are: [None, 'micro', 'macro', 'weighted']"
             )
 
         if not isinstance(beta, float):
@@ -191,13 +209,15 @@ class FBetaScore(tf.keras.metrics.Metric):
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
 class F1Score(FBetaScore):
-    """Computes F-1 Score.
+    r"""Computes F-1 Score.
 
     It is the harmonic mean of precision and recall.
-    Output range is [0, 1]. Works for both multi-class
+    Output range is `[0, 1]`. Works for both multi-class
     and multi-label classification.
 
-    F-1 = 2 * (precision * recall) / (precision + recall)
+    $$
+    F_1 = 2 \cdot \frac{\textrm{precision} \cdot \textrm{recall}}{\textrm{precision} + \textrm{recall}}
+    $$
 
     Args:
         num_classes: Number of unique classes in the dataset.
@@ -207,13 +227,15 @@ class F1Score(FBetaScore):
         threshold: Elements of `y_pred` above threshold are
             considered to be 1, and the rest 0. If threshold is
             None, the argmax is converted to 1, and the rest 0.
+        name: (Optional) String name of the metric instance.
+        dtype: (Optional) Data type of the metric result.
 
     Returns:
-        F-1 Score: float
+        F-1 Score: float.
 
     Raises:
         ValueError: If the `average` has values other than
-        [None, micro, macro, weighted].
+        [None, 'micro', 'macro', 'weighted'].
 
     `average` parameter behavior:
         None: Scores for each class are returned
@@ -228,6 +250,20 @@ class F1Score(FBetaScore):
         weighted: Metrics are computed for each class
             and returns the mean weighted by the
             number of true instances in each class.
+
+    Usage:
+
+    >>> metric = tfa.metrics.F1Score(num_classes=3, threshold=0.5)
+    >>> y_true = np.array([[1, 1, 1],
+    ...                    [1, 0, 0],
+    ...                    [1, 1, 0]], np.int32)
+    >>> y_pred = np.array([[0.2, 0.6, 0.7],
+    ...                    [0.2, 0.6, 0.6],
+    ...                    [0.6, 0.8, 0.0]], np.float32)
+    >>> metric.update_state(y_true, y_pred)
+    >>> result = metric.result()
+    >>> result.numpy()
+    array([0.5      , 0.8      , 0.6666667], dtype=float32)
     """
 
     @typechecked
diff --git a/tensorflow_addons/metrics/geometric_mean.py b/tensorflow_addons/metrics/geometric_mean.py
index bce8e15..a5caef9 100644
--- a/tensorflow_addons/metrics/geometric_mean.py
+++ b/tensorflow_addons/metrics/geometric_mean.py
@@ -31,15 +31,18 @@ class GeometricMean(Metric):
     that uses the sum of values, it uses the product of the values to
     represent typical values in a set of numbers.
 
-    Note: `tfa.metrics.GeometricMean` can be used the same as `tf.keras.metrics.Mean`
+    Note: `tfa.metrics.GeometricMean` can be used the same as `tf.keras.metrics.Mean`.
+
+    Args:
+        name: (Optional) String name of the metric instance.
+        dtype: (Optional) Data type of the metric result.
 
     Usage:
 
-    >>> m = tfa.metrics.GeometricMean()
-    >>> m.update_state([1, 3, 5, 7, 9])
-    >>> m.result().numpy()
+    >>> metric = tfa.metrics.GeometricMean()
+    >>> metric.update_state([1, 3, 5, 7, 9])
+    >>> metric.result().numpy()
     3.9362833
-
     """
 
     @typechecked
diff --git a/tensorflow_addons/metrics/hamming.py b/tensorflow_addons/metrics/hamming.py
index 9702b03..5e08ecb 100644
--- a/tensorflow_addons/metrics/hamming.py
+++ b/tensorflow_addons/metrics/hamming.py
@@ -30,22 +30,19 @@ def hamming_distance(actuals: TensorLike, predictions: TensorLike) -> tf.Tensor:
     are different.
 
     Args:
-        actuals: actual target value
-        predictions: predicted value
+        actuals: actual target value.
+        predictions: predicted value.
 
     Returns:
-        hamming distance: float
+        hamming distance: float.
 
     Usage:
 
-    ```python
-    actuals = tf.constant([1, 1, 0, 0, 1, 0, 1, 0, 0, 1],
-                          dtype=tf.int32)
-    predictions = tf.constant([1, 0, 0, 0, 1, 0, 0, 1, 0, 1],
-                              dtype=tf.int32)
-    result = hamming_distance(actuals, predictions)
-    print('Hamming distance: ', result.numpy())
-    ```
+    >>> y_true = np.array([1, 1, 0, 0, 1, 0, 1, 0, 0, 1], dtype=np.int32)
+    >>> y_pred = np.array([1, 0, 0, 0, 1, 0, 0, 1, 0, 1], dtype=np.int32)
+    >>> hamming_distance(y_true, y_pred).numpy()
+    0.3
+
     """
     result = tf.not_equal(actuals, predictions)
     not_eq = tf.reduce_sum(tf.cast(result, tf.float32))
@@ -65,50 +62,20 @@ def hamming_loss_fn(
     of labels.
 
     In multi-class classification, hamming loss is calculated as the
-    hamming distance between `actual` and `predictions`.
+    hamming distance between `y_true` and `y_pred`.
     In multi-label classification, hamming loss penalizes only the
     individual labels.
 
     Args:
-
-        y_true: actual target value
-        y_pred: predicted target value
+        y_true: actual target value.
+        y_pred: predicted target value.
         threshold: Elements of `y_pred` greater than threshold are
             converted to be 1, and the rest 0. If threshold is
             None, the argmax is converted to 1, and the rest 0.
-        mode: multi-class or multi-label
+        mode: multi-class or multi-label.
 
     Returns:
-
-        hamming loss: float
-
-    Usage:
-
-    ```python
-    # multi-class hamming loss
-    hl = HammingLoss(mode='multiclass', threshold=0.6)
-    actuals = tf.constant([[1, 0, 0, 0],[0, 0, 1, 0],
-                       [0, 0, 0, 1],[0, 1, 0, 0]],
-                      dtype=tf.float32)
-    predictions = tf.constant([[0.8, 0.1, 0.1, 0],
-                               [0.2, 0, 0.8, 0],
-                               [0.05, 0.05, 0.1, 0.8],
-                               [1, 0, 0, 0]],
-                          dtype=tf.float32)
-    hl.update_state(actuals, predictions)
-    print('Hamming loss: ', hl.result().numpy()) # 0.25
-
-    # multi-label hamming loss
-    hl = HammingLoss(mode='multilabel', threshold=0.8)
-    actuals = tf.constant([[1, 0, 1, 0],[0, 1, 0, 1],
-                       [0, 0, 0,1]], dtype=tf.int32)
-    predictions = tf.constant([[0.82, 0.5, 0.90, 0],
-                               [0, 1, 0.4, 0.98],
-                               [0.89, 0.79, 0, 0.3]],
-                               dtype=tf.float32)
-    hl.update_state(actuals, predictions)
-    print('Hamming loss: ', hl.result().numpy()) # 0.16666667
-    ```
+        hamming loss: float.
     """
     if mode not in ["multiclass", "multilabel"]:
         raise TypeError("mode must be either multiclass or multilabel]")
@@ -134,7 +101,53 @@ def hamming_loss_fn(
 
 
 class HammingLoss(MeanMetricWrapper):
-    """Computes hamming loss."""
+    """Computes hamming loss.
+
+    Hamming loss is the fraction of wrong labels to the total number
+    of labels.
+
+    In multi-class classification, hamming loss is calculated as the
+    hamming distance between `y_true` and `y_pred`.
+    In multi-label classification, hamming loss penalizes only the
+    individual labels.
+
+    Args:
+        threshold: Elements of `y_pred` greater than threshold are
+            converted to be 1, and the rest 0. If threshold is
+            None, the argmax is converted to 1, and the rest 0.
+        mode: multi-class or multi-label.
+        name: (Optional) string name of the metric instance.
+        dtype: (Optional) data type of the metric result.
+
+    Usage:
+
+    >>> # multi-class hamming loss
+    >>> metric = HammingLoss(mode='multiclass', threshold=0.6)
+    >>> y_true = np.array([[1.0, 0.0, 0.0, 0.0],
+    ...                    [0.0, 0.0, 1.0, 0.0],
+    ...                    [0.0, 0.0, 0.0, 1.0],
+    ...                    [0.0, 1.0, 0.0, 0.0]], dtype=np.float32)
+    >>> y_pred = np.array([[0.8, 0.1, 0.1, 0.0],
+    ...                    [0.2, 0.0, 0.8, 0.0],
+    ...                    [0.05, 0.05, 0.1, 0.8],
+    ...                    [1.0, 0.0, 0.0, 0.0]], dtype=np.float32)
+    >>> metric.update_state(y_true, y_pred)
+    <tf.Variable 'UnreadVariable' shape=() dtype=float32, numpy=4.0>
+    >>> metric.result().numpy()
+    0.25
+    >>> # multi-label hamming loss
+    >>> metric = HammingLoss(mode='multilabel', threshold=0.8)
+    >>> y_true = np.array([[1, 0, 1, 0],
+    ...                    [0, 1, 0, 1],
+    ...                    [0, 0, 0, 1]], dtype=np.int32)
+    >>> y_pred = np.array([[0.82, 0.5, 0.90, 0],
+    ...                    [0, 1, 0.4, 0.98],
+    ...                    [0.89, 0.79, 0, 0.3]], dtype=np.float32)
+    >>> metric.update_state(y_true, y_pred)
+    <tf.Variable 'UnreadVariable' shape=() dtype=float32, numpy=3.0>
+    >>> metric.result().numpy()
+    0.16666667
+    """
 
     @typechecked
     def __init__(
diff --git a/tensorflow_addons/metrics/matthews_correlation_coefficient.py b/tensorflow_addons/metrics/matthews_correlation_coefficient.py
index e6afa53..0d5a4b9 100644
--- a/tensorflow_addons/metrics/matthews_correlation_coefficient.py
+++ b/tensorflow_addons/metrics/matthews_correlation_coefficient.py
@@ -41,19 +41,20 @@ class MatthewsCorrelationCoefficient(tf.keras.metrics.Metric):
     MCC = (TP * TN) - (FP * FN) /
           ((TP + FP) * (TP + FN) * (TN + FP ) * (TN + FN))^(1/2)
 
+    Args:
+        num_classes : Number of unique classes in the dataset.
+        name: (Optional) String name of the metric instance.
+        dtype: (Optional) Data type of the metric result.
+
     Usage:
-    ```python
-    actuals = tf.constant([[1.0], [1.0], [1.0], [0.0]],
-             dtype=tf.float32)
-    preds = tf.constant([[1.0], [0.0], [1.0], [1.0]],
-             dtype=tf.float32)
-    # Matthews correlation coefficient
-    mcc = MatthewsCorrelationCoefficient(num_classes=1)
-    mcc.update_state(actuals, preds)
-    print('Matthews correlation coefficient is:',
-    mcc.result().numpy())
-    # Matthews correlation coefficient is : -0.33333334
-    ```
+
+    >>> y_true = np.array([[1.0], [1.0], [1.0], [0.0]], dtype=np.float32)
+    >>> y_pred = np.array([[1.0], [0.0], [1.0], [1.0]], dtype=np.float32)
+    >>> metric = tfa.metrics.MatthewsCorrelationCoefficient(num_classes=1)
+    >>> metric.update_state(y_true, y_pred)
+    >>> result = metric.result()
+    >>> result.numpy()
+    array([-0.33333334], dtype=float32)
     """
 
     @typechecked
@@ -64,14 +65,7 @@ class MatthewsCorrelationCoefficient(tf.keras.metrics.Metric):
         dtype: AcceptableDTypes = None,
         **kwargs
     ):
-        """Creates a Matthews Correlation Coefficient instance.
-
-        Args:
-            num_classes : Number of unique classes in the dataset.
-            name: (Optional) String name of the metric instance.
-            dtype: (Optional) Data type of the metric result.
-            Defaults to `tf.float32`.
-        """
+        """Creates a Matthews Correlation Coefficient instance."""
         super().__init__(name=name, dtype=dtype)
         self.num_classes = num_classes
         self.true_positives = self.add_weight(
diff --git a/tensorflow_addons/metrics/multilabel_confusion_matrix.py b/tensorflow_addons/metrics/multilabel_confusion_matrix.py
index d13a355..93afc68 100644
--- a/tensorflow_addons/metrics/multilabel_confusion_matrix.py
+++ b/tensorflow_addons/metrics/multilabel_confusion_matrix.py
@@ -37,39 +37,53 @@ class MultiLabelConfusionMatrix(Metric):
     Consider classification problem with two classes
     (i.e num_classes=2).
 
-    Resultant matrix `M` will be in the shape of (num_classes, 2, 2).
-
-    Every class `i` has a dedicated 2*2 matrix that contains:
-
-    - true negatives for class i in M(0,0)
-    - false positives for class i in M(0,1)
-    - false negatives for class i in M(1,0)
-    - true positives for class i in M(1,1)
-
-    ```python
-    # multilabel confusion matrix
-    y_true = tf.constant([[1, 0, 1], [0, 1, 0]],
-             dtype=tf.int32)
-    y_pred = tf.constant([[1, 0, 0],[0, 1, 1]],
-             dtype=tf.int32)
-    output = MultiLabelConfusionMatrix(num_classes=3)
-    output.update_state(y_true, y_pred)
-    print('Confusion matrix:', output.result().numpy())
-
-    # Confusion matrix: [[[1 0] [0 1]] [[1 0] [0 1]]
-                      [[0 1] [1 0]]]
-
-    # if multiclass input is provided
-    y_true = tf.constant([[1, 0, 0], [0, 1, 0]],
-             dtype=tf.int32)
-    y_pred = tf.constant([[1, 0, 0],[0, 0, 1]],
-             dtype=tf.int32)
-    output = MultiLabelConfusionMatrix(num_classes=3)
-    output.update_state(y_true, y_pred)
-    print('Confusion matrix:', output.result().numpy())
-
-    # Confusion matrix: [[[1 0] [0 1]] [[1 0] [1 0]] [[1 1] [0 0]]]
-    ```
+    Resultant matrix `M` will be in the shape of `(num_classes, 2, 2)`.
+
+    Every class `i` has a dedicated matrix of shape `(2, 2)` that contains:
+
+    - true negatives for class `i` in `M(0,0)`
+    - false positives for class `i` in `M(0,1)`
+    - false negatives for class `i` in `M(1,0)`
+    - true positives for class `i` in `M(1,1)`
+
+    Args:
+        num_classes: `int`, the number of labels the prediction task can have.
+        name: (Optional) string name of the metric instance.
+        dtype: (Optional) data type of the metric result.
+
+    Usage:
+
+    >>> # multilabel confusion matrix
+    >>> y_true = np.array([[1, 0, 1], [0, 1, 0]], dtype=np.int32)
+    >>> y_pred = np.array([[1, 0, 0], [0, 1, 1]], dtype=np.int32)
+    >>> metric = tfa.metrics.MultiLabelConfusionMatrix(num_classes=3)
+    >>> metric.update_state(y_true, y_pred)
+    >>> result = metric.result()
+    >>> result.numpy()  #doctest: -DONT_ACCEPT_BLANKLINE
+    array([[[1., 0.],
+            [0., 1.]],
+    <BLANKLINE>
+           [[1., 0.],
+            [0., 1.]],
+    <BLANKLINE>
+           [[0., 1.],
+            [1., 0.]]], dtype=float32)
+    >>> # if multiclass input is provided
+    >>> y_true = np.array([[1, 0, 0], [0, 1, 0]], dtype=np.int32)
+    >>> y_pred = np.array([[1, 0, 0], [0, 0, 1]], dtype=np.int32)
+    >>> metric = tfa.metrics.MultiLabelConfusionMatrix(num_classes=3)
+    >>> metric.update_state(y_true, y_pred)
+    >>> result = metric.result()
+    >>> result.numpy() #doctest: -DONT_ACCEPT_BLANKLINE
+    array([[[1., 0.],
+            [0., 1.]],
+    <BLANKLINE>
+           [[1., 0.],
+            [1., 0.]],
+    <BLANKLINE>
+           [[1., 1.],
+            [0., 0.]]], dtype=float32)
+
     """
 
     @typechecked
diff --git a/tensorflow_addons/metrics/r_square.py b/tensorflow_addons/metrics/r_square.py
index 04c3866..c1a9bd2 100644
--- a/tensorflow_addons/metrics/r_square.py
+++ b/tensorflow_addons/metrics/r_square.py
@@ -44,29 +44,36 @@ def _reduce_average(
 class RSquare(Metric):
     """Compute R^2 score.
 
-     This is also called the [coefficient of determination
-     ](https://en.wikipedia.org/wiki/Coefficient_of_determination).
-     It tells how close are data to the fitted regression line.
-
-     - Highest score can be 1.0 and it indicates that the predictors
-       perfectly accounts for variation in the target.
-     - Score 0.0 indicates that the predictors do not
-       account for variation in the target.
-     - It can also be negative if the model is worse.
-
-     The sample weighting for this metric implementation mimics the
-     behaviour of the [scikit-learn implementation
-     ](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html)
-     of the same metric.
-
-     Usage:
-     ```python
-     actuals = tf.constant([1, 4, 3], dtype=tf.float32)
-     preds = tf.constant([2, 4, 4], dtype=tf.float32)
-     result = tf.keras.metrics.RSquare()
-     result.update_state(actuals, preds)
-     print('R^2 score is: ', r1.result().numpy()) # 0.57142866
-    ```
+    This is also called the [coefficient of determination
+    ](https://en.wikipedia.org/wiki/Coefficient_of_determination).
+    It tells how close are data to the fitted regression line.
+
+    - Highest score can be 1.0 and it indicates that the predictors
+        perfectly accounts for variation in the target.
+    - Score 0.0 indicates that the predictors do not
+        account for variation in the target.
+    - It can also be negative if the model is worse.
+
+    The sample weighting for this metric implementation mimics the
+    behaviour of the [scikit-learn implementation
+    ](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html)
+    of the same metric.
+
+    Args:
+        multioutput: `string`, the reduce method for scores.
+            Should be one of `["raw_values", "uniform_average", "variance_weighted"]`.
+        name: (Optional) string name of the metric instance.
+        dtype: (Optional) data type of the metric result.
+
+    Usage:
+
+    >>> y_true = np.array([1, 4, 3], dtype=np.float32)
+    >>> y_pred = np.array([2, 4, 4], dtype=np.float32)
+    >>> metric = tfa.metrics.r_square.RSquare()
+    >>> metric.update_state(y_true, y_pred)
+    >>> result = metric.result()
+    >>> result.numpy()
+    0.57142854
     """
 
     @typechecked
diff --git a/tensorflow_addons/text/__init__.py b/tensorflow_addons/text/__init__.py
index d9c8383..dfb6379 100644
--- a/tensorflow_addons/text/__init__.py
+++ b/tensorflow_addons/text/__init__.py
@@ -16,9 +16,11 @@
 
 # Conditional Random Field
 from tensorflow_addons.text.crf import crf_binary_score
+from tensorflow_addons.text.crf import crf_constrained_decode
 from tensorflow_addons.text.crf import crf_decode
 from tensorflow_addons.text.crf import crf_decode_backward
 from tensorflow_addons.text.crf import crf_decode_forward
+from tensorflow_addons.text.crf import crf_filtered_inputs
 from tensorflow_addons.text.crf import crf_forward
 from tensorflow_addons.text.crf import crf_log_likelihood
 from tensorflow_addons.text.crf import crf_log_norm
diff --git a/tensorflow_addons/text/crf.py b/tensorflow_addons/text/crf.py
index 2707895..8bc9901 100644
--- a/tensorflow_addons/text/crf.py
+++ b/tensorflow_addons/text/crf.py
@@ -24,6 +24,32 @@ from typing import Optional
 # https://github.com/tensorflow/tensorflow/issues/29075 is resolved
 
 
+def crf_filtered_inputs(inputs: TensorLike, tag_bitmap: TensorLike) -> tf.Tensor:
+    """Constrains the inputs to filter out certain tags at each time step.
+
+    tag_bitmap limits the allowed tags at each input time step.
+    This is useful when an observed output at a given time step needs to be
+    constrained to a selected set of tags.
+
+    Args:
+      inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
+          to use as input to the CRF layer.
+      tag_bitmap: A [batch_size, max_seq_len, num_tags] boolean tensor
+          representing all active tags at each index for which to calculate the
+          unnormalized score.
+    Returns:
+      filtered_inputs: A [batch_size] vector of unnormalized sequence scores.
+    """
+
+    # set scores of filtered out inputs to be -inf.
+    filtered_inputs = tf.where(
+        tag_bitmap,
+        inputs,
+        tf.fill(tf.shape(inputs), tf.cast(float("-inf"), inputs.dtype)),
+    )
+    return filtered_inputs
+
+
 def crf_sequence_score(
     inputs: TensorLike,
     tag_indices: TensorLike,
@@ -107,11 +133,7 @@ def crf_multitag_sequence_score(
     """
     tag_bitmap = tf.cast(tag_bitmap, dtype=tf.bool)
     sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
-    filtered_inputs = tf.where(
-        tag_bitmap,
-        inputs,
-        tf.fill(tf.shape(inputs), tf.cast(float("-inf"), inputs.dtype)),
-    )
+    filtered_inputs = crf_filtered_inputs(inputs, tag_bitmap)
 
     # If max_seq_len is 1, we skip the score calculation and simply gather the
     # unary potentials of all active tags.
@@ -559,3 +581,32 @@ def crf_decode(
         return tf.cond(
             tf.equal(tf.shape(potentials)[1], 1), _single_seq_fn, _multi_seq_fn
         )
+
+
+def crf_constrained_decode(
+    potentials: TensorLike,
+    tag_bitmap: TensorLike,
+    transition_params: TensorLike,
+    sequence_length: TensorLike,
+) -> tf.Tensor:
+    """Decode the highest scoring sequence of tags under constraints.
+
+    This is a function for tensor.
+
+    Args:
+      potentials: A [batch_size, max_seq_len, num_tags] tensor of
+                unary potentials.
+      tag_bitmap: A [batch_size, max_seq_len, num_tags] boolean tensor
+          representing all active tags at each index for which to calculate the
+          unnormalized score.
+      transition_params: A [num_tags, num_tags] matrix of
+                binary potentials.
+      sequence_length: A [batch_size] vector of true sequence lengths.
+    Returns:
+      decode_tags: A [batch_size, max_seq_len] matrix, with dtype `tf.int32`.
+                  Contains the highest scoring tag indices.
+      best_score: A [batch_size] vector, containing the score of `decode_tags`.
+    """
+
+    filtered_potentials = crf_filtered_inputs(potentials, tag_bitmap)
+    return crf_decode(filtered_potentials, transition_params, sequence_length)
diff --git a/tools/docker/sanity_check.Dockerfile b/tools/docker/sanity_check.Dockerfile
index 68393a9..9300da5 100644
--- a/tools/docker/sanity_check.Dockerfile
+++ b/tools/docker/sanity_check.Dockerfile
@@ -94,7 +94,7 @@ RUN apt-get update && apt-get install -y rsync
 COPY ./ /addons
 WORKDIR /addons
 RUN pip install --no-deps -e .
-RUN python docs/build_docs.py
+RUN python tools/docs/build_docs.py
 RUN touch /ok.txt
 
 # -------------------------------
diff --git a/docs/BUILD b/tools/docs/BUILD
similarity index 100%
rename from docs/BUILD
rename to tools/docs/BUILD
diff --git a/tools/docs/Readme.md b/tools/docs/Readme.md
new file mode 100644
index 0000000..7c574e5
--- /dev/null
+++ b/tools/docs/Readme.md
@@ -0,0 +1,19 @@
+## 1. Generated API docs
+
+[tensorflow.org/addons/api_docs/python/tfa](https://tensorflow.org/addons/api_docs/python/tfa)
+
+`build_docs.py` controls executed this docs generation. To test-run it:
+
+```bash
+# Install dependencies:
+pip install -r tools/install_deps/doc_requirements.txt
+
+# Build tool:
+bazel build //tools/docs:build_docs
+
+# Generate API doc:
+# Use current branch
+bazel-bin/tools/docs/build_docs --git_branch=$(git rev-parse --abbrev-ref HEAD)
+# or specified explicitly
+bazel-bin/tools/docs/build_docs --git_branch=master --output_dir=/tmp/tfa_api
+```


Test patch:
diff --git a/tensorflow_addons/text/tests/crf_test.py b/tensorflow_addons/text/tests/crf_test.py
index fc92a54..d199d19 100644
--- a/tensorflow_addons/text/tests/crf_test.py
+++ b/tensorflow_addons/text/tests/crf_test.py
@@ -35,6 +35,82 @@ def calculate_sequence_score(inputs, transition_params, tag_indices, sequence_le
     return expected_unary_score + expected_binary_score
 
 
+def brute_force_decode(sequence_lengths, inputs, transition_params):
+    num_words = inputs.shape[0]
+    num_tags = inputs.shape[1]
+
+    all_sequence_scores = []
+    all_sequences = []
+
+    tag_indices_iterator = itertools.product(range(num_tags), repeat=sequence_lengths)
+    inputs = tf.expand_dims(inputs, 0)
+    sequence_lengths = tf.expand_dims(sequence_lengths, 0)
+    transition_params = tf.constant(transition_params)
+
+    # Compare the dynamic program with brute force computation.
+    for tag_indices in tag_indices_iterator:
+        tag_indices = list(tag_indices)
+        tag_indices.extend([0] * (num_words - sequence_lengths))
+        all_sequences.append(tag_indices)
+        sequence_score = text.crf_sequence_score(
+            inputs=inputs,
+            tag_indices=tf.expand_dims(tag_indices, 0),
+            sequence_lengths=sequence_lengths,
+            transition_params=transition_params,
+        )
+        sequence_score = tf.squeeze(sequence_score, [0])
+        all_sequence_scores.append(sequence_score)
+
+    expected_max_sequence_index = np.argmax(all_sequence_scores)
+    expected_max_sequence = all_sequences[expected_max_sequence_index]
+    expected_max_score = all_sequence_scores[expected_max_sequence_index]
+    return expected_max_sequence, expected_max_score
+
+
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_filtered_inputs(dtype):
+    # Test both the length-1 and regular cases.
+    sequence_lengths_list = [np.array(3, dtype=np.int32), np.array(1, dtype=np.int32)]
+    inputs_list = [
+        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype),
+        np.array([[4, 5, -3]], dtype=dtype),
+    ]
+    tag_bitmap_list = [
+        np.array(
+            [
+                [True, False, False],
+                [False, True, True],
+                [False, True, True],
+                [False, True, True],
+            ],
+            dtype=np.bool,
+        ),
+        np.array([[False, True, True]], dtype=np.bool),
+    ]
+    neg_inf = float("-inf")
+    expected_filtered_inputs_list = [
+        np.array(
+            [[4, neg_inf, neg_inf], [neg_inf, -1, 3], [neg_inf, 2, 1], [neg_inf, 0, 0]],
+            dtype=dtype,
+        ),
+        np.array([[neg_inf, 5, -3]], dtype=dtype),
+    ]
+    for sequence_lengths, inputs, tag_bitmap, expected_filtered_inputs in zip(
+        sequence_lengths_list,
+        inputs_list,
+        tag_bitmap_list,
+        expected_filtered_inputs_list,
+    ):
+        filtered_inputs = text.crf_filtered_inputs(
+            inputs=tf.expand_dims(inputs, 0), tag_bitmap=tf.expand_dims(tag_bitmap, 0)
+        )
+        filtered_inputs = tf.squeeze(filtered_inputs, [0])
+
+        test_utils.assert_allclose_according_to_type(
+            filtered_inputs, expected_filtered_inputs
+        )
+
+
 @pytest.mark.parametrize("dtype", [np.float16, np.float32])
 def test_crf_sequence_score(dtype):
     transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
@@ -309,29 +385,9 @@ def test_crf_decode(dtype):
     for sequence_lengths, inputs, tag_indices in zip(
         sequence_lengths_list, inputs_list, tag_indices_list
     ):
-        num_words = inputs.shape[0]
-        num_tags = inputs.shape[1]
-
-        all_sequence_scores = []
-        all_sequences = []
-
-        # Compare the dynamic program with brute force computation.
-        for tag_indices in itertools.product(range(num_tags), repeat=sequence_lengths):
-            tag_indices = list(tag_indices)
-            tag_indices.extend([0] * (num_words - sequence_lengths))
-            all_sequences.append(tag_indices)
-            sequence_score = text.crf_sequence_score(
-                inputs=tf.expand_dims(inputs, 0),
-                tag_indices=tf.expand_dims(tag_indices, 0),
-                sequence_lengths=tf.expand_dims(sequence_lengths, 0),
-                transition_params=tf.constant(transition_params),
-            )
-            sequence_score = tf.squeeze(sequence_score, [0])
-            all_sequence_scores.append(sequence_score)
-
-        expected_max_sequence_index = np.argmax(all_sequence_scores)
-        expected_max_sequence = all_sequences[expected_max_sequence_index]
-        expected_max_score = all_sequence_scores[expected_max_sequence_index]
+        expected_max_sequence, expected_max_score = brute_force_decode(
+            sequence_lengths, inputs, transition_params
+        )
 
         actual_max_sequence, actual_max_score = text.crf_decode(
             tf.expand_dims(inputs, 0),
@@ -350,6 +406,61 @@ def test_crf_decode(dtype):
         )
 
 
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_constrained_decode(dtype):
+    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
+    # Test both the length-1 and regular cases.
+    sequence_lengths_list = [np.array(3, dtype=np.int32), np.array(1, dtype=np.int32)]
+    inputs_list = [
+        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype),
+        np.array([[4, 5, -3]], dtype=dtype),
+    ]
+    tag_bitmap_list = [
+        np.array(
+            [
+                [True, False, False],
+                [False, True, True],
+                [False, True, True],
+                [False, True, True],
+            ],
+            dtype=np.bool,
+        ),
+        np.array([[False, True, True]], dtype=np.bool),
+    ]
+    for sequence_lengths, inputs, tag_bitmap in zip(
+        sequence_lengths_list, inputs_list, tag_bitmap_list
+    ):
+        filtered_inputs = text.crf_filtered_inputs(
+            inputs=tf.expand_dims(inputs, 0), tag_bitmap=tf.expand_dims(tag_bitmap, 0)
+        )
+
+        expected_max_sequence, expected_max_score = text.crf_decode(
+            filtered_inputs,
+            tf.constant(transition_params),
+            tf.expand_dims(sequence_lengths, 0),
+        )
+
+        expected_max_sequence = tf.squeeze(expected_max_sequence, [0])
+        expected_max_score = tf.squeeze(expected_max_score, [0])
+
+        actual_max_sequence, actual_max_score = text.crf_constrained_decode(
+            tf.expand_dims(inputs, 0),
+            tf.expand_dims(tag_bitmap, 0),
+            tf.constant(transition_params),
+            tf.expand_dims(sequence_lengths, 0),
+        )
+
+        actual_max_sequence = tf.squeeze(actual_max_sequence, [0])
+        actual_max_score = tf.squeeze(actual_max_score, [0])
+
+        test_utils.assert_allclose_according_to_type(
+            actual_max_score, expected_max_score, 1e-6, 1e-6
+        )
+        assert list(actual_max_sequence[:sequence_lengths]) == list(
+            expected_max_sequence[:sequence_lengths]
+        )
+
+
 def test_crf_decode_zero_seq_length():
     """Test that crf_decode works when sequence_length contains one or more
     zeros."""