Repository: tensorflow/addons

Issue:
Add constrained decoding to CRF
**Describe the feature and the current behavior/state.**
We can reuse code from the `crf_multitag_sequence_score` to allow constrained decoding. 

**Relevant information**
- Are you willing to contribute it (yes/no): yes
- Are you willing to maintain it going forward? (yes/no): yes
- Is there a relevant academic paper? (if so, where): we only need to reuse the functions in `tfa.text.crf` to get this functionality
- Is there already an implementation in another framework? (if so, where): yes easy to extend `tfa.text.crf`
- Was it part of tf.contrib? (if so, where): 

**Which API type would this fall under (layer, metric, optimizer, etc.)**
This will be under `tfa.text.crf`

**Who will benefit with this feature?**
People in NLP community, constrained decoding is important to deal with padding issues as well to denote start and end of a sequence. 

**Any other info.**

Discussion:

Patch:
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index a016b7a..bde99ae 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -105,6 +105,10 @@
 /tensorflow_addons/layers/tests/esn_test.py @pedrolarben
 /tensorflow_addons/layers/snake.py @failure-to-thrive
 /tensorflow_addons/layers/tests/snake_test.py @failure-to-thrive
+/tensorflow_addons/layers/stochastic_depth.py @mhstadler @windqaq
+/tensorflow_addons/layers/tests/stochastic_depth_test.py @mhstadler @windqaq
+/tensorflow_addons/layers/noisy_dense.py @leonshams
+/tensorflow_addons/layers/tests/noisy_dense_test.py @leonshams
 
 /tensorflow_addons/losses/contrastive.py @windqaq
 /tensorflow_addons/losses/tests/contrastive_test.py @windqaq
diff --git a/tensorflow_addons/layers/__init__.py b/tensorflow_addons/layers/__init__.py
index 594d025..7c89f95 100644
--- a/tensorflow_addons/layers/__init__.py
+++ b/tensorflow_addons/layers/__init__.py
@@ -38,3 +38,5 @@ from tensorflow_addons.layers.spatial_pyramid_pooling import SpatialPyramidPooli
 from tensorflow_addons.layers.tlu import TLU
 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
diff --git a/tensorflow_addons/layers/noisy_dense.py b/tensorflow_addons/layers/noisy_dense.py
new file mode 100644
index 0000000..647b28d
--- /dev/null
+++ b/tensorflow_addons/layers/noisy_dense.py
@@ -0,0 +1,264 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import tensorflow as tf
+from tensorflow.keras import (
+    activations,
+    initializers,
+    regularizers,
+    constraints,
+)
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import InputSpec
+from typeguard import typechecked
+
+from tensorflow_addons.utils import types
+
+
+def _scale_noise(x):
+    return tf.sign(x) * tf.sqrt(tf.abs(x))
+
+
+@tf.keras.utils.register_keras_serializable(package="Addons")
+class NoisyDense(tf.keras.layers.Dense):
+    r"""Noisy dense layer that injects random noise to the weights of dense layer.
+
+    Noisy dense layers are fully connected layers whose weights and biases are
+    augmented by factorised Gaussian noise. The factorised Gaussian noise is
+    controlled through gradient descent by a second weights layer.
+
+    A `NoisyDense` layer implements the operation:
+    $$
+    \mathrm{NoisyDense}(x) =
+    \mathrm{activation}(\mathrm{dot}(x, \mu + (\sigma \cdot \epsilon))
+    + \mathrm{bias})
+    $$
+    where $\mu$ is the standard weights layer, $\epsilon$ is the factorised
+    Gaussian noise, and $\sigma$ is a second weights layer which controls
+    $\epsilon$.
+
+    Note: bias only added if `use_bias` is `True`.
+
+    Example:
+
+    >>> # Create a `Sequential` model and add a NoisyDense
+    >>> # layer as the first layer.
+    >>> model = tf.keras.models.Sequential()
+    >>> model.add(tf.keras.Input(shape=(16,)))
+    >>> model.add(NoisyDense(32, activation='relu'))
+    >>> # Now the model will take as input arrays of shape (None, 16)
+    >>> # and output arrays of shape (None, 32).
+    >>> # Note that after the first layer, you don't need to specify
+    >>> # the size of the input anymore:
+    >>> model.add(NoisyDense(32))
+    >>> model.output_shape
+    (None, 32)
+
+    Arguments:
+      units: Positive integer, dimensionality of the output space.
+      sigma: A float between 0-1 used as a standard deviation figure and is
+        applied to the gaussian noise layer (`sigma_kernel` and `sigma_bias`).
+      activation: Activation function to use.
+        If you don't specify anything, no activation is applied
+        (ie. "linear" activation: `a(x) = x`).
+      use_bias: Boolean, whether the layer uses a bias vector.
+      kernel_regularizer: Regularizer function applied to
+        the `kernel` weights matrix.
+      bias_regularizer: Regularizer function applied to the bias vector.
+      activity_regularizer: Regularizer function applied to
+        the output of the layer (its "activation").
+      kernel_constraint: Constraint function applied to
+        the `kernel` weights matrix.
+      bias_constraint: Constraint function applied to the bias vector.
+
+    Input shape:
+      N-D tensor with shape: `(batch_size, ..., input_dim)`.
+      The most common situation would be
+      a 2D input with shape `(batch_size, input_dim)`.
+
+    Output shape:
+      N-D tensor with shape: `(batch_size, ..., units)`.
+      For instance, for a 2D input with shape `(batch_size, input_dim)`,
+      the output would have shape `(batch_size, units)`.
+
+    References:
+      - [Noisy Networks for Explanation](https://arxiv.org/pdf/1706.10295.pdf)
+    """
+
+    @typechecked
+    def __init__(
+        self,
+        units: int,
+        sigma: float = 0.5,
+        activation: types.Activation = None,
+        use_bias: bool = True,
+        kernel_regularizer: types.Regularizer = None,
+        bias_regularizer: types.Regularizer = None,
+        activity_regularizer: types.Regularizer = None,
+        kernel_constraint: types.Constraint = None,
+        bias_constraint: types.Constraint = None,
+        **kwargs
+    ):
+        super().__init__(
+            units=units,
+            activation=activation,
+            use_bias=use_bias,
+            kernel_regularizer=kernel_regularizer,
+            bias_regularizer=bias_regularizer,
+            activity_regularizer=activity_regularizer,
+            kernel_constraint=kernel_constraint,
+            bias_constraint=bias_constraint,
+            **kwargs,
+        )
+        delattr(self, "kernel_initializer")
+        delattr(self, "bias_initializer")
+        self.sigma = sigma
+
+    def build(self, input_shape):
+        # Make sure dtype is correct
+        dtype = tf.dtypes.as_dtype(self.dtype or K.floatx())
+        if not (dtype.is_floating or dtype.is_complex):
+            raise TypeError(
+                "Unable to build `Dense` layer with non-floating point "
+                "dtype %s" % (dtype,)
+            )
+
+        input_shape = tf.TensorShape(input_shape)
+        self.last_dim = tf.compat.dimension_value(input_shape[-1])
+        sqrt_dim = self.last_dim ** (1 / 2)
+        if self.last_dim is None:
+            raise ValueError(
+                "The last dimension of the inputs to `Dense` "
+                "should be defined. Found `None`."
+            )
+        self.input_spec = InputSpec(min_ndim=2, axes={-1: self.last_dim})
+
+        sigma_init = initializers.Constant(value=self.sigma / sqrt_dim)
+        mu_init = initializers.RandomUniform(minval=-1 / sqrt_dim, maxval=1 / sqrt_dim)
+
+        # Learnable parameters
+        self.sigma_kernel = self.add_weight(
+            "sigma_kernel",
+            shape=[self.last_dim, self.units],
+            initializer=sigma_init,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+            dtype=self.dtype,
+            trainable=True,
+        )
+
+        self.mu_kernel = self.add_weight(
+            "mu_kernel",
+            shape=[self.last_dim, self.units],
+            initializer=mu_init,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+            dtype=self.dtype,
+            trainable=True,
+        )
+
+        if self.use_bias:
+            self.sigma_bias = self.add_weight(
+                "sigma_bias",
+                shape=[
+                    self.units,
+                ],
+                initializer=sigma_init,
+                regularizer=self.bias_regularizer,
+                constraint=self.bias_constraint,
+                dtype=self.dtype,
+                trainable=True,
+            )
+
+            self.mu_bias = self.add_weight(
+                "mu_bias",
+                shape=[
+                    self.units,
+                ],
+                initializer=mu_init,
+                regularizer=self.bias_regularizer,
+                constraint=self.bias_constraint,
+                dtype=self.dtype,
+                trainable=True,
+            )
+        else:
+            self.sigma_bias = None
+            self.mu_bias = None
+        self._reset_noise()
+        self.built = True
+
+    @property
+    def kernel(self):
+        return self.mu_kernel + (self.sigma_kernel * self.eps_kernel)
+
+    @property
+    def bias(self):
+        if self.use_bias:
+            return self.mu_bias + (self.sigma_bias * self.eps_bias)
+
+    def _reset_noise(self):
+        """Create the factorised Gaussian noise."""
+
+        dtype = self._compute_dtype_object
+
+        # Generate random noise
+        eps_i = tf.random.normal([self.last_dim, self.units], dtype=dtype)
+        eps_j = tf.random.normal(
+            [
+                self.units,
+            ],
+            dtype=dtype,
+        )
+
+        # Scale the random noise
+        self.eps_kernel = _scale_noise(eps_i) * _scale_noise(eps_j)
+        self.eps_bias = _scale_noise(eps_j)
+
+    def _remove_noise(self):
+        """Remove the factorised Gaussian noise."""
+
+        dtype = self._compute_dtype_object
+        self.eps_kernel = tf.zeros([self.last_dim, self.units], dtype=dtype)
+        self.eps_bias = tf.zeros([self.units], dtype=dtype)
+
+    def call(self, inputs, reset_noise=True, remove_noise=False):
+        # Generate fixed parameters added as the noise
+        if remove_noise:
+            self._remove_noise()
+        elif reset_noise:
+            self._reset_noise()
+
+        # TODO(WindQAQ): Replace this with `dense()` once public.
+        return super().call(inputs)
+
+    def get_config(self):
+        # TODO(WindQAQ): Get rid of this hacky way.
+        config = super(tf.keras.layers.Dense, self).get_config()
+        config.update(
+            {
+                "units": self.units,
+                "sigma": self.sigma,
+                "activation": activations.serialize(self.activation),
+                "use_bias": self.use_bias,
+                "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),
+            }
+        )
+        return config
diff --git a/tensorflow_addons/layers/stochastic_depth.py b/tensorflow_addons/layers/stochastic_depth.py
new file mode 100644
index 0000000..3cf0df4
--- /dev/null
+++ b/tensorflow_addons/layers/stochastic_depth.py
@@ -0,0 +1,88 @@
+import tensorflow as tf
+from typeguard import typechecked
+
+
+@tf.keras.utils.register_keras_serializable(package="Addons")
+class StochasticDepth(tf.keras.layers.Layer):
+    """Stochastic Depth layer.
+
+    Implements Stochastic Depth as described in
+    [Deep Networks with Stochastic Depth](https://arxiv.org/abs/1603.09382), to randomly drop residual branches
+    in residual architectures.
+
+    Usage:
+    Residual architectures with fixed depth, use residual branches that are merged back into the main network
+    by adding the residual branch back to the input:
+
+    >>> input = np.ones((1, 3, 3, 1), dtype = np.float32)
+    >>> residual = tf.keras.layers.Conv2D(1, 1)(input)
+    >>> output = tf.keras.layers.Add()([input, residual])
+    >>> output.shape
+    TensorShape([1, 3, 3, 1])
+
+    StochasticDepth acts as a drop-in replacement for the addition:
+
+    >>> input = np.ones((1, 3, 3, 1), dtype = np.float32)
+    >>> residual = tf.keras.layers.Conv2D(1, 1)(input)
+    >>> output = tfa.layers.StochasticDepth()([input, residual])
+    >>> output.shape
+    TensorShape([1, 3, 3, 1])
+
+    At train time, StochasticDepth returns:
+
+    $$
+    x[0] + b_l * x[1],
+    $$
+
+    where $b_l$ is a random Bernoulli variable with probability $P(b_l = 1) = p_l$
+
+    At test time, StochasticDepth rescales the activations of the residual branch based on the survival probability ($p_l$):
+
+    $$
+    x[0] + p_l * x[1]
+    $$
+
+    Arguments:
+        survival_probability: float, the probability of the residual branch being kept.
+
+    Call Arguments:
+        inputs:  List of `[shortcut, residual]` where `shortcut`, and `residual` are tensors of equal shape.
+
+    Output shape:
+        Equal to the shape of inputs `shortcut`, and `residual`
+    """
+
+    @typechecked
+    def __init__(self, survival_probability: float = 0.5, **kwargs):
+        super().__init__(**kwargs)
+
+        self.survival_probability = survival_probability
+
+    def call(self, x, training=None):
+        if not isinstance(x, list) or len(x) != 2:
+            raise ValueError("input must be a list of length 2.")
+
+        shortcut, residual = x
+
+        # Random bernoulli variable indicating whether the branch should be kept or not or not
+        b_l = tf.keras.backend.random_bernoulli([], p=self.survival_probability)
+
+        def _call_train():
+            return shortcut + b_l * residual
+
+        def _call_test():
+            return shortcut + self.survival_probability * residual
+
+        return tf.keras.backend.in_train_phase(
+            _call_train, _call_test, training=training
+        )
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0]
+
+    def get_config(self):
+        base_config = super().get_config()
+
+        config = {"survival_probability": self.survival_probability}
+
+        return {**base_config, **config}
diff --git a/tensorflow_addons/rnn/esn_cell.py b/tensorflow_addons/rnn/esn_cell.py
index d36be9f..e85f638 100644
--- a/tensorflow_addons/rnn/esn_cell.py
+++ b/tensorflow_addons/rnn/esn_cell.py
@@ -32,6 +32,18 @@ class ESNCell(keras.layers.AbstractRNNCell):
         "The "echo state" approach to analysing and training recurrent neural networks".
         GMD Report148, German National Research Center for Information Technology, 2001.
         https://www.researchgate.net/publication/215385037
+
+    Example:
+
+    >>> inputs = np.random.random([30,23,9]).astype(np.float32)
+    >>> ESNCell = tfa.rnn.ESNCell(4)
+    >>> rnn = tf.keras.layers.RNN(ESNCell, return_sequences=True, return_state=True)
+    >>> outputs, memory_state = rnn(inputs)
+    >>> outputs.shape
+    TensorShape([30, 23, 4])
+    >>> memory_state.shape
+    TensorShape([30, 4])
+
     Arguments:
         units: Positive integer, dimensionality in the reservoir.
         connectivity: Float between 0 and 1.
diff --git a/tensorflow_addons/rnn/layer_norm_lstm_cell.py b/tensorflow_addons/rnn/layer_norm_lstm_cell.py
index 3a3a359..adff626 100644
--- a/tensorflow_addons/rnn/layer_norm_lstm_cell.py
+++ b/tensorflow_addons/rnn/layer_norm_lstm_cell.py
@@ -46,6 +46,19 @@ class LayerNormLSTMCell(keras.layers.LSTMCell):
 
     "Recurrent Dropout without Memory Loss"
     Stanislau Semeniuta, Aliaksei Severyn, Erhardt Barth.
+
+    Example:
+
+    >>> inputs = np.random.random([30,23,9]).astype(np.float32)
+    >>> lnLSTMCell = tfa.rnn.LayerNormLSTMCell(4)
+    >>> rnn = tf.keras.layers.RNN(lnLSTMCell, return_sequences=True, return_state=True)
+    >>> outputs, memory_state, carry_state = rnn(inputs)
+    >>> outputs.shape
+    TensorShape([30, 23, 4])
+    >>> memory_state.shape
+    TensorShape([30, 4])
+    >>> carry_state.shape
+    TensorShape([30, 4])
     """
 
     @typechecked
diff --git a/tensorflow_addons/rnn/layer_norm_simple_rnn_cell.py b/tensorflow_addons/rnn/layer_norm_simple_rnn_cell.py
index f9562c6..cc5c2b6 100644
--- a/tensorflow_addons/rnn/layer_norm_simple_rnn_cell.py
+++ b/tensorflow_addons/rnn/layer_norm_simple_rnn_cell.py
@@ -37,6 +37,17 @@ class LayerNormSimpleRNNCell(keras.layers.SimpleRNNCell):
         "Layer Normalization." ArXiv:1607.06450 [Cs, Stat],
         July 21, 2016. http://arxiv.org/abs/1607.06450
 
+    Example:
+
+    >>> inputs = np.random.random([30,23,9]).astype(np.float32)
+    >>> lnsRNNCell = tfa.rnn.LayerNormSimpleRNNCell(4)
+    >>> rnn = tf.keras.layers.RNN(lnsRNNCell, return_sequences=True, return_state=True)
+    >>> outputs, memory_state = rnn(inputs)
+    >>> outputs.shape
+    TensorShape([30, 23, 4])
+    >>> memory_state.shape
+    TensorShape([30, 4])
+
     Arguments:
       units: Positive integer, dimensionality of the output space.
       activation: Activation function to use.
@@ -89,25 +100,19 @@ class LayerNormSimpleRNNCell(keras.layers.SimpleRNNCell):
 
     Examples:
 
-    ```python
-    import numpy as np
-    import tensorflow.keras as keras
-    import tensorflow_addons as tfa
-
-    inputs = np.random.random([32, 10, 8]).astype(np.float32)
-    rnn = keras.layers.RNN(tfa.rnn.LayerNormSimpleRNNCell(4))
-
-    output = rnn(inputs)  # The output has shape `[32, 4]`.
-
-    rnn = keras.layers.RNN(
-        tfa.rnn.LayerNormSimpleRNNCell(4),
-        return_sequences=True,
-        return_state=True)
+    >>> inputs = np.random.random([32, 10, 8]).astype(np.float32)
+    >>> rnn = tf.keras.layers.RNN(tfa.rnn.LayerNormSimpleRNNCell(4))
+    >>> output = rnn(inputs)  # The output has shape `[32, 4]`.
+    >>> rnn = tf.keras.layers.RNN(
+    ... tfa.rnn.LayerNormSimpleRNNCell(4),
+    ... return_sequences=True,
+    ... return_state=True)
+    >>> whole_sequence_output, final_state = rnn(inputs)
+    >>> whole_sequence_output
+    <tf.Tensor: shape=(32, 10, 4), dtype=float32, numpy=...>
+    >>> final_state
+    <tf.Tensor: shape=(32, 4), dtype=float32, numpy=...>
 
-    # whole_sequence_output has shape `[32, 10, 4]`.
-    # final_state has shape `[32, 4]`.
-    whole_sequence_output, final_state = rnn(inputs)
-    ```
     """
 
     @typechecked
diff --git a/tensorflow_addons/rnn/nas_cell.py b/tensorflow_addons/rnn/nas_cell.py
index 05054a0..6b6686b 100644
--- a/tensorflow_addons/rnn/nas_cell.py
+++ b/tensorflow_addons/rnn/nas_cell.py
@@ -38,6 +38,19 @@ class NASCell(keras.layers.AbstractRNNCell):
     "Neural Architecture Search with Reinforcement Learning" Proc. ICLR 2017.
 
     The class uses an optional projection layer.
+
+    Example:
+
+    >>> inputs = np.random.random([30,23,9]).astype(np.float32)
+    >>> NASCell = tfa.rnn.NASCell(4)
+    >>> rnn = tf.keras.layers.RNN(NASCell, return_sequences=True, return_state=True)
+    >>> outputs, memory_state, carry_state = rnn(inputs)
+    >>> outputs.shape
+    TensorShape([30, 23, 4])
+    >>> memory_state.shape
+    TensorShape([30, 4])
+    >>> carry_state.shape
+    TensorShape([30, 4])
     """
 
     # NAS cell's architecture base.
diff --git a/tensorflow_addons/rnn/peephole_lstm_cell.py b/tensorflow_addons/rnn/peephole_lstm_cell.py
index 1f791d5..658db08 100644
--- a/tensorflow_addons/rnn/peephole_lstm_cell.py
+++ b/tensorflow_addons/rnn/peephole_lstm_cell.py
@@ -39,14 +39,16 @@ class PeepholeLSTMCell(tf.keras.layers.LSTMCell):
 
     Example:
 
-    ```python
-    # Create 2 PeepholeLSTMCells
-    peephole_lstm_cells = [PeepholeLSTMCell(size) for size in [128, 256]]
-    # Create a layer composed sequentially of the peephole LSTM cells.
-    layer = RNN(peephole_lstm_cells)
-    input = keras.Input((timesteps, input_dim))
-    output = layer(input)
-    ```
+    >>> inputs = np.random.random([30,23,9]).astype(np.float32)
+    >>> LSTMCell = tfa.rnn.PeepholeLSTMCell(4)
+    >>> rnn = tf.keras.layers.RNN(LSTMCell, return_sequences=True, return_state=True)
+    >>> outputs, memory_state, carry_state = rnn(inputs)
+    >>> outputs.shape
+    TensorShape([30, 23, 4])
+    >>> memory_state.shape
+    TensorShape([30, 4])
+    >>> carry_state.shape
+    TensorShape([30, 4])
     """
 
     def build(self, input_shape):
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)


Test patch:
diff --git a/tensorflow_addons/layers/tests/noisy_dense_test.py b/tensorflow_addons/layers/tests/noisy_dense_test.py
new file mode 100644
index 0000000..9f76307
--- /dev/null
+++ b/tensorflow_addons/layers/tests/noisy_dense_test.py
@@ -0,0 +1,141 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests NoisyDense layer."""
+
+
+import pytest
+import numpy as np
+
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras.mixed_precision.experimental import Policy
+
+from tensorflow_addons.utils import test_utils
+from tensorflow_addons.layers.noisy_dense import NoisyDense
+
+
+@pytest.mark.parametrize(
+    "input_shape", [(3, 2), (3, 4, 2), (None, None, 2), (3, 4, 5, 2)]
+)
+def test_noisy_dense(input_shape):
+    test_utils.layer_test(NoisyDense, kwargs={"units": 3}, input_shape=input_shape)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
+def test_noisy_dense_dtype(dtype):
+    inputs = tf.convert_to_tensor(
+        np.random.randint(low=0, high=7, size=(2, 2)), dtype=dtype
+    )
+    layer = NoisyDense(5, dtype=dtype, name="noisy_dense_" + dtype)
+    outputs = layer(inputs)
+    np.testing.assert_array_equal(outputs.dtype, dtype)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_noisy_dense_with_policy():
+    inputs = tf.convert_to_tensor(np.random.randint(low=0, high=7, size=(2, 2)))
+    layer = NoisyDense(5, dtype=Policy("mixed_float16"), name="noisy_dense_policy")
+    outputs = layer(inputs)
+    output_signature = layer.compute_output_signature(
+        tf.TensorSpec(dtype="float16", shape=(2, 2))
+    )
+    np.testing.assert_array_equal(output_signature.dtype, tf.dtypes.float16)
+    np.testing.assert_array_equal(output_signature.shape, (2, 5))
+    np.testing.assert_array_equal(outputs.dtype, "float16")
+    np.testing.assert_array_equal(layer.mu_kernel.dtype, "float32")
+    np.testing.assert_array_equal(layer.sigma_kernel.dtype, "float32")
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_noisy_dense_regularization():
+    layer = NoisyDense(
+        3,
+        kernel_regularizer=keras.regularizers.l1(0.01),
+        bias_regularizer="l1",
+        activity_regularizer="l2",
+        name="noisy_dense_reg",
+    )
+    layer(keras.backend.variable(np.ones((2, 4))))
+    np.testing.assert_array_equal(5, len(layer.losses))
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_noisy_dense_constraints():
+    k_constraint = keras.constraints.max_norm(0.01)
+    b_constraint = keras.constraints.max_norm(0.01)
+    layer = NoisyDense(
+        3,
+        kernel_constraint=k_constraint,
+        bias_constraint=b_constraint,
+        name="noisy_dense_constriants",
+    )
+    layer(keras.backend.variable(np.ones((2, 4))))
+    np.testing.assert_array_equal(layer.mu_kernel.constraint, k_constraint)
+    np.testing.assert_array_equal(layer.sigma_kernel.constraint, k_constraint)
+    np.testing.assert_array_equal(layer.mu_bias.constraint, b_constraint)
+    np.testing.assert_array_equal(layer.sigma_bias.constraint, b_constraint)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_noisy_dense_automatic_reset_noise():
+    inputs = tf.convert_to_tensor(np.random.randint(low=0, high=7, size=(2, 2)))
+    layer = NoisyDense(5, name="noise_dense_auto_reset_noise")
+    layer(inputs)
+    initial_eps_kernel = layer.eps_kernel
+    initial_eps_bias = layer.eps_bias
+    layer(inputs)
+    new_eps_kernel = layer.eps_kernel
+    new_eps_bias = layer.eps_bias
+    np.testing.assert_raises(
+        AssertionError,
+        np.testing.assert_array_equal,
+        initial_eps_kernel,
+        new_eps_kernel,
+    )
+    np.testing.assert_raises(
+        AssertionError,
+        np.testing.assert_array_equal,
+        initial_eps_bias,
+        new_eps_bias,
+    )
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_noisy_dense_remove_noise():
+    inputs = tf.convert_to_tensor(np.random.randint(low=0, high=7, size=(2, 2)))
+    layer = NoisyDense(5, name="noise_dense_manual_reset_noise")
+    layer(inputs)
+    initial_eps_kernel = layer.eps_kernel
+    initial_eps_bias = layer.eps_bias
+    layer(inputs, reset_noise=False, remove_noise=True)
+    new_eps_kernel = layer.eps_kernel
+    new_eps_bias = layer.eps_bias
+    kernel_zeros = tf.zeros(initial_eps_kernel.shape, dtype=initial_eps_kernel.dtype)
+    bias_zeros = tf.zeros(initial_eps_bias.shape, dtype=initial_eps_kernel.dtype)
+    np.testing.assert_raises(
+        AssertionError,
+        np.testing.assert_array_equal,
+        initial_eps_kernel,
+        new_eps_kernel,
+    )
+    np.testing.assert_raises(
+        AssertionError,
+        np.testing.assert_array_equal,
+        initial_eps_bias,
+        new_eps_bias,
+    )
+    np.testing.assert_array_equal(kernel_zeros, new_eps_kernel)
+    np.testing.assert_array_equal(bias_zeros, new_eps_bias)
diff --git a/tensorflow_addons/layers/tests/stochastic_depth_test.py b/tensorflow_addons/layers/tests/stochastic_depth_test.py
new file mode 100644
index 0000000..1122016
--- /dev/null
+++ b/tensorflow_addons/layers/tests/stochastic_depth_test.py
@@ -0,0 +1,58 @@
+import pytest
+import numpy as np
+import tensorflow as tf
+
+from tensorflow_addons.layers.stochastic_depth import StochasticDepth
+from tensorflow_addons.utils import test_utils
+
+_KEEP_SEED = 1111
+_DROP_SEED = 2222
+
+
+@pytest.mark.parametrize("seed", [_KEEP_SEED, _DROP_SEED])
+@pytest.mark.parametrize("training", [True, False])
+def stochastic_depth_test(seed, training):
+    np.random.seed(seed)
+    tf.random.set_seed(seed)
+
+    survival_probability = 0.5
+
+    shortcut = np.asarray([[0.2, 0.1, 0.4]]).astype(np.float32)
+    residual = np.asarray([[0.2, 0.4, 0.5]]).astype(np.float32)
+
+    if training:
+        if seed == _KEEP_SEED:
+            # shortcut + residual
+            expected_output = np.asarray([[0.4, 0.5, 0.9]]).astype(np.float32)
+        elif seed == _DROP_SEED:
+            # shortcut
+            expected_output = np.asarray([[0.2, 0.1, 0.4]]).astype(np.float32)
+    else:
+        # shortcut + p_l * residual
+        expected_output = np.asarray([[0.3, 0.3, 0.65]]).astype(np.float32)
+
+    test_utils.layer_test(
+        StochasticDepth,
+        kwargs={"survival_probability": survival_probability},
+        input_data=[shortcut, residual],
+        expected_output=expected_output,
+    )
+
+
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
+def test_with_mixed_precision_policy():
+    policy = tf.keras.mixed_precision.experimental.global_policy()
+
+    shortcut = np.asarray([[0.2, 0.1, 0.4]])
+    residual = np.asarray([[0.2, 0.4, 0.5]])
+
+    output = StochasticDepth()([shortcut, residual])
+
+    assert output.dtype == policy.compute_dtype
+
+
+def test_serialization():
+    stoch_depth = StochasticDepth(survival_probability=0.5)
+    serialized_stoch_depth = tf.keras.layers.serialize(stoch_depth)
+    new_layer = tf.keras.layers.deserialize(serialized_stoch_depth)
+    assert stoch_depth.get_config() == new_layer.get_config()
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."""