Input channels C is not equal to kernel channels * group when BN layer's axis=-2

[maybeLee]

Describe the bug
I was constructing a three-layer keras model using the following script

import keras
# This is the buggy model built from scratch
layer_stack = [
  keras.layers.DepthwiseConv2D(kernel_size=(3,3), strides=(2,2)),
  keras.layers.BatchNormalization(axis=-2),  # if we set axis=-1, the bug will not be manifested
  keras.layers.ReLU(max_value=6)
]
keras_input = keras.layers.Input((6,6,512))
layer_input = keras_input
for layer in layer_stack:
  y = layer(layer_input)
  layer_input = y
new_model = keras.models.Model(keras_input, y)
new_model.summary()
new_model.save("generated_model.h5")

This model can be successfully constructed and do model prediction using keras. But when I convert it to onnx model and making inference with it, it will raise a runtime error.

import tensorflow as tf
import tf2onnx
def transform_onnx(model_path: str):
    model = keras.models.load_model(model_path)
    model_name = model_path.split(".h5")[0]
    onnx_path = f"{model_name}.onnx"
    input_shape = model.layers[0].input_shape[0]
    spec = (tf.TensorSpec(input_shape, tf.float32, name="input"),)
    _, _ = tf2onnx.convert.from_keras(model, input_signature=spec, \
        opset=15, output_path=onnx_path)
transform_onnx("generated_model.h5")

import onnxruntime as ort
onnx_path = "generated_model.onnx"
onnx_model = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
input_name = onnx_model.get_inputs()[0].name
output_name = onnx_model.get_outputs()[0].name
input = input.astype('float32')
onnx_pred = onnx_model.run([output_name], {input_name: input})[0]

Error Message:

---------------------------------------------------------------------------
Fail                                      Traceback (most recent call last)
[<ipython-input-27-2745ef9ee442>](https://8h48y5r9dcw-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20221128-060047-RC00_491313853#) in <module>
      5 output_name = onnx_model.get_outputs()[0].name
      6 input = input.astype('float32')
----> 7 onnx_pred = onnx_model.run([output_name], {input_name: input})[0]

[/usr/local/lib/python3.7/dist-packages/onnxruntime/capi/onnxruntime_inference_collection.py](https://8h48y5r9dcw-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20221128-060047-RC00_491313853#) in run(self, output_names, input_feed, run_options)
    198             output_names = [output.name for output in self._outputs_meta]
    199         try:
--> 200             return self._sess.run(output_names, input_feed, run_options)
    201         except C.EPFail as err:
    202             if self._enable_fallback:

Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Conv node. Name:'model_5/depthwise_conv2d_4/depthwise' Status Message: Input channels C is not equal to kernel channels * group. C: 512 kernel channels: 2 group: 512

Looks like it may be caused by the incorrect inference of kernel channels and group. Note that if I set the axis of BN to -1, this bug will not be manifested.

Urgency

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 18.04*):
• TensorFlow Version:2.9.2
• Python version:3.7
• ONNX version (if applicable, e.g. 1.11*):1.13.1
• ONNXRuntime version (if applicable, e.g. 1.11*):1.13.1

To Reproduce

You can reproduce this issue through the colab link:
https://colab.research.google.com/drive/1Cwdeauu7hDJJxwmzxPzULobvQfuxeDUU?usp=sharing

Screenshots

Additional context