# Buggy code
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)


# Traceback (ValueError)
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].

# Fixed code
def build(self, input_shape):
    super(LayerNormLSTMCell, self).build(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])