Repository: tensorflow/addons

Issue:
LayerNormLSTMCellTest failed in the latest tf2-nightly version
**System information**
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow):
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
- TensorFlow installed from (source or binary):
- TensorFlow version (use command below):
- TensorFlow Addons installed from (source, PyPi):
- TensorFlow Addons version:
- Python version and type (eg. Anaconda Python, Stock Python as in Mac, or homebrew installed Python etc):
- Bazel version (if compiling from source):
- GCC/Compiler version (if compiling from source):
- Is GPU used? (yes/no):
- GPU model (if used):

**Describe the bug**

Please check the CI testing log:
https://source.cloud.google.com/results/invocations/c2c4cc8c-41a5-4eaa-87e9-bf513a92e0c0/targets/tensorflow_addons%2Fubuntu%2Fcpu%2Fpy2%2Fcontinuous/log

```python
ERROR: testCellOutput (__main__.LayerNormLSTMCellTest)
testCellOutput (__main__.LayerNormLSTMCellTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/absl/third_party/unittest3_backport/case.py", line 37, in testPartExecutor
    yield
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/absl/third_party/unittest3_backport/case.py", line 162, in run
    testMethod()
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/framework/test_util.py", line 1053, in decorated
    f(self, *args, **kwargs)
  File "/tmpfs/tmp/bazel/sandbox/linux-sandbox/8/execroot/__main__/bazel-out/k8-opt/bin/tensorflow_addons/rnn/cell_test.runfiles/__main__/tensorflow_addons/rnn/cell_test.py", line 199, in testCellOutput
    output, output_states = cell(x, state)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 716, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/keras/layers/recurrent.py", line 135, in call
    inputs, states = cell.call(inputs, states, **kwargs)
  File "/tmpfs/tmp/bazel/sandbox/linux-sandbox/8/execroot/__main__/bazel-out/k8-opt/bin/tensorflow_addons/rnn/cell_test.runfiles/__main__/tensorflow_addons/rnn/cell.py", line 338, in call
    z = self.kernel_norm(keras.backend.dot(inputs, self.kernel))
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 716, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/keras/layers/normalization.py", line 1022, in call
    variance_epsilon=self.epsilon)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/ops/nn_impl.py", line 1395, in batch_normalization
    return x * math_ops.cast(inv, x.dtype) + math_ops.cast(
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/ops/math_ops.py", line 890, in binary_op_wrapper
    return func(x, y, name=name)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/ops/math_ops.py", line 1197, in _mul_dispatch
    return gen_math_ops.mul(x, y, name=name)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py", line 6592, in mul
    "Mul", x=x, y=y, name=name)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/framework/op_def_library.py", line 793, in _apply_op_helper
    op_def=op_def)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/framework/ops.py", line 3279, in create_op
    attrs, op_def, compute_device)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/framework/ops.py", line 3348, in _create_op_internal
    op_def=op_def)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/framework/ops.py", line 1716, in __init__
    control_input_ops)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow_core/python/framework/ops.py", line 1556, in _create_c_op
    raise ValueError(str(e))
ValueError: Dimensions must be equal, but are 8 and 2 for 'stacked_rnn_cells/kernel_norm/batchnorm/mul_1' (op: 'Mul') with input shapes: [1,8], [1,2].

```


Discussion:
Quite busy this week... but cursory look at the commits on 05-25 seems like this may be the one affecting this cell? 
https://github.com/tensorflow/tensorflow/commit/eb2ba2256d1a4d0625f54c55f077c41b1629dfb9

cc @qlzh727




tensorflow/tensorflow@eb2ba22 seems to be harmless since the layer is expected to have built = True after invoking the build. We might have some bug in the original cell implementation.

I am quite busy in this week and next few, will take a look if I have time. I don't want to addon to be pinged on a stale version. Feel free to disable the test with a TODO for now.

Resolving patch:
diff --git a/tensorflow_addons/rnn/cell.py b/tensorflow_addons/rnn/cell.py
--- a/tensorflow_addons/rnn/cell.py
+++ b/tensorflow_addons/rnn/cell.py
@@ -321,10 +321,9 @@ def __init__(self,
 
     def build(self, input_shape):
         super(LayerNormLSTMCell, self).build(input_shape)
-        norm_input_shape = [input_shape[0], self.units]
-        self.kernel_norm.build(norm_input_shape)
-        self.recurrent_norm.build(norm_input_shape)
-        self.state_norm.build(norm_input_shape)
+        self.kernel_norm.build([input_shape[0], self.units * 4])
+        self.recurrent_norm.build([input_shape[0], self.units * 4])
+        self.state_norm.build([input_shape[0], self.units])
 
     def call(self, inputs, states, training=None):
         h_tm1 = states[0]  # previous memory state


Test patch: