Invalid reduction dimension (2 for input with 2 dimension(s) [Op:Sum]

[MilkFiish]

This error occurs when I do dynamic execution work with keras.layers.UnitNormalization, but not when I do static inference.
The code below is executed in the tensorflow backend environment.

import os
import re
import tensorflow as tf
import numpy as np
os.environ['KERAS_BACKEND']='tensorflow'
import keras

layer = keras.layers.UnitNormalization(
    axis=[ 1, 2 ],
    trainable=True,
    autocast=True,
)

result_static = layer.compute_output_shape([2, 3])

result_dynamic = layer(
    inputs=np.random.rand(*[2, 3]),
)

When I print result_static, it's [2, 3]. However, the dynamic part causes the error

File ...\lib\site-packages\tensorflow\python\eager\execute.py:53, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     51 try:
     52   ctx.ensure_initialized()
---> 53   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     54                                       inputs, attrs, num_outputs)
     55 except core._NotOkStatusException as e:
     56   if name is not None:

InvalidArgumentError: Exception encountered when calling UnitNormalization.call().

{{function_node __wrapped__Sum_device_/job:localhost/replica:0/task:0/device:CPU:0}} Invalid reduction dimension (2 for input with 2 dimension(s) [Op:Sum]

Arguments received by UnitNormalization.call():
  • inputs=tf.Tensor(shape=(2, 3), dtype=float32)

By traceback information, it seems that an illegal input has been sent to the backend, and an input check may need to be added.

[sanskarmodi8]

Hi @MilkFiish,

I looked into your issue and found out that the error is due to a mismatch between the dimensionality of the input tensor and the axis specified in the UnitNormalization layer.
The axis=[1, 2] argument is expecting a 3D tensor, but the input you are providing is a 2D tensor (shape=(2, 3)), which doesn't have the second dimension specified in the axis.

You should use np.random.rand(2, 3, 1) instead

[MilkFiish]

If there is an input error, I thought it should have an error report for both static and dynamic execution, which is why I bring up the issue.
Hopefully, the team could add checks for such incorrect input, rather than reporting errors while dynamically executing to the backend and confusing users.

[sanskarmodi8]

You're correct. I have created a Pull Request regarding the same by adding validation checks #20237

[sachinprasadhs]

@MilkFiish , Now with the validation check in the above linked PR, it is throwing the proper error as ValueError: Axis [1, 2] is out of bounds for input shape [2, 3], attaching the Gist reproduced using Keras-Nightly