ONEDNN - Layer Normalization error

[Chizkiyahu]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras):
• OS Platform and Distribution : ubuntu 18, ubuntu 20,redhat 7
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.9.
• Python version: 3.8.11
• Bazel version (if compiling from source):
• GPU model and memory:
• Exact command to reproduce: see code below

Describe the problem.

Describe the problem clearly here. Be sure to convey here why it's a bug in Keras or why the requested feature is needed.

Describe the current behavior. the code fail here but I am assuming the source in another place

Describe the expected behavior. code pass without fail like if TF_ENABLE_ONEDNN_OPTS set to 0

• Do you want to contribute a PR? (yes/no):no
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):

Standalone code to reproduce the issue. run on computer with intel dl booster (without GPU)

import os
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '1'
import tensorflow as tf
import numpy as np

input_layer = tf.keras.layers.Input(shape=(3,24), dtype=np.float32)
data = tf.quantization.fake_quant_with_min_max_vars(inputs=input_layer, min=0, max=3.5, num_bits=8)
data = tf.keras.layers.LayerNormalization(axis=[-2,-1], center=True,scale=True)(data)
data = tf.quantization.fake_quant_with_min_max_vars(inputs=data, min=0, max=3.5, num_bits=8)
model = tf.keras.Model(inputs=input_layer, outputs=data)
model.compile()
model.summary()

data = np.random.rand(1,3,24)
res = model.predict(data)

Source code / logs.

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/path/to/python/python3.8/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/path/to/python/python3.8/site-packages/tensorflow/python/eager/execute.py", line 54, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:

Detected at node 'model/layer_normalization/add' defined at (most recent call last):
    File "<stdin>", line 1, in <module>
    File "/path/to/python/python3.8/site-packages/keras/utils/traceback_utils.py", line 64, in error_handler
      return fn(*args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/engine/training.py", line 2033, in predict
      tmp_batch_outputs = self.predict_function(iterator)
    File "/path/to/python/python3.8/site-packages/keras/engine/training.py", line 1845, in predict_function
      return step_function(self, iterator)
    File "/path/to/python/python3.8/site-packages/keras/engine/training.py", line 1834, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/path/to/python/python3.8/site-packages/keras/engine/training.py", line 1823, in run_step
      outputs = model.predict_step(data)
    File "/path/to/python/python3.8/site-packages/keras/engine/training.py", line 1791, in predict_step
      return self(x, training=False)
    File "/path/to/python/python3.8/site-packages/keras/utils/traceback_utils.py", line 64, in error_handler
      return fn(*args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/engine/training.py", line 490, in __call__
      return super().__call__(*args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/utils/traceback_utils.py", line 64, in error_handler
      return fn(*args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/engine/base_layer.py", line 1014, in __call__
      outputs = call_fn(inputs, *args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/utils/traceback_utils.py", line 92, in error_handler
      return fn(*args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/engine/functional.py", line 458, in call
      return self._run_internal_graph(
    File "/path/to/python/python3.8/site-packages/keras/engine/functional.py", line 596, in _run_internal_graph
      outputs = node.layer(*args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/utils/traceback_utils.py", line 64, in error_handler
      return fn(*args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/engine/base_layer.py", line 1014, in __call__
      outputs = call_fn(inputs, *args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/utils/traceback_utils.py", line 92, in error_handler
      return fn(*args, **kwargs)
    File "/path/to/python/python3.8/site-packages/keras/layers/normalization/layer_normalization.py", line 324, in call
      outputs = outputs + tf.cast(offset, outputs.dtype)
Node: 'model/layer_normalization/add'
scale must be 1D tensor[1,3,24]
	 [[{{node model/layer_normalization/add}}]] [Op:__inference_predict_function_217]