LayerNormalization fails when given tuple as axis input

[HarikrishnanBalagopal]

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux-4.14.137+-x86_64-with-Ubuntu-18.04-bionic
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: NA
• TensorFlow installed from (source or binary): Google colab has it preinstalled
• TensorFlow version (use command below): 1.14.0
• Python version: (major, minor, micro, releaselevel, serial) (3, 6, 8, 'final', 0)
• Bazel version (if compiling from source): NA
• GCC/Compiler version (if compiling from source): NA
• CUDA/cuDNN version: CUDA Version: 10.1
• GPU model and memory: Tesla K80 11441MiB

You can collect some of this information using our environment capture
script
You can also obtain the TensorFlow version with: 1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)" 2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior

Sequential([Input(shape=(64, 64, 3), dtype=np.float32), LayerNormalization(axis=(-3, -2, -1))])

fails with error

---------------------------------------------------------------------------

TypeError                                 Traceback (most recent call last)

<ipython-input-22-18267fdf3265> in <module>()
----> 1 Sequential([Input(shape=(64, 64, 3), dtype=np.float32), LayerNormalization(axis=(-3, -2, -1))])

6 frames

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/normalization.py in build(self, input_shape)
    957     for idx, x in enumerate(self.axis):
    958       if x < 0:
--> 959         self.axis[idx] = ndims + x
    960 
    961     # Validate axes

TypeError: 'tuple' object does not support item assignment

This is because the lines:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/normalization.py#L929-L930

    if isinstance(axis, (list, tuple)):
      self.axis = axis[:]

make a copy of the tuple instead of converting it to a list and later the lines:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/normalization.py#L954-L956

    # Convert axis to list and resolve negatives
    if isinstance(self.axis, int):
      self.axis = [self.axis]

do not take care of the case when axis is a tuple.

Describe the expected behavior
Fix is simple, replace the lines:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/normalization.py#L954-L956
with:

    # Convert axis to list and resolve negatives
    if isinstance(self.axis, int):
      self.axis = [self.axis]
    elif isinstance(self.axis, tuple):
      self.axis = list(axis)

Code to reproduce the issue
Provide a reproducible test case that is the bare minimum necessary to generate the problem.

Sequential([Input(shape=(64, 64, 3), dtype=np.float32), LayerNormalization(axis=(-3, -2, -1))])

Other info / logs
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

[HarikrishnanBalagopal]

The same code appears in Tensorflow 2 as well:

https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/python/keras/layers/normalization.py#L929-L930

    if isinstance(axis, (list, tuple)):
      self.axis = axis[:]

https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/python/keras/layers/normalization.py#L954-L956

    # Convert axis to list and resolve negatives
    if isinstance(self.axis, int):
      self.axis = [self.axis]

[HarikrishnanBalagopal]

What is the convention? Some functions accept tuples, lists and ints for the axis argument. Others only accept lists and integers.

[jvishnuvardhan]

I can reproduce the issue even with tensorflow==1.15.0rc0. However, If you change axis from tuple to list, then there is no error. The error is also clearly points the source of error TypeError: 'tuple' object does not support item assignment.

Please check the gist with 1.15.0rc0`. Thanks!

[HarikrishnanBalagopal]

Should I submit a pull request with the fix?

[robieta]

Please feel free to create a fix and tag me.

[HarikrishnanBalagopal]

@robieta I did post a fix in the original post. I didn't create a pull request because I am not sure what the convention is. Some functions in tensorflow accept tuples, lists and ints as axis while others only accept list and ints for axis.

Fix is simple, replace the lines:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/normalization.py#L954-L956
with:

 # Convert axis to list and resolve negatives
    if isinstance(self.axis, int):
      self.axis = [self.axis]
    elif isinstance(self.axis, tuple):
      self.axis = list(axis)

[HarikrishnanBalagopal]

@robieta Also note that the same code is present in tensorflow 2. Maybe tag the issue as tf2 as well.

[robieta]

I think the fix you proposed is reasonable. Particularly for higher level symbols like Layers we try to be fairly permissive as long as it's unambiguous what the user wants. (Which in this case it is.)

[HarikrishnanBalagopal]

@robieta I created a pull request #32463 .
Maybe add a test case for this?

[HarikrishnanBalagopal]

@robieta So the pull request got merged into version 1.14, but the same code is present in r2.0.
https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/python/keras/layers/normalization.py#L954-L956
Can you merge into that as well? or do I need to create a separate pull request for that?

[robieta]

I don't think this warrants a cherrypick given how far along the 2.0 RC process is. It will automatically be picked up in 2.1.

[HarikrishnanBalagopal]

Closing issue since pull request fixing it has been merged in v1.14

[tensorflow-bot]

Are you satisfied with the resolution of your issue?
Yes
No