onnx max_pool_with_argmax op does not return the same indices than tensorflow

[mpaillassa]

Describe the bug
When running the tensorflow max_pool_with_argmax op with onnxruntime, I find that the op returns does not return the same indices than tensorflow does.

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
• Tensorflow Version: 2.4.1
• Python version: 3.8.5
• tf2onnx: from source (1.9.0)

To Reproduce
To generate a model showing the bug:

import tensorflow as tf
import numpy as np

class Bug(tf.keras.Model):   
    def __init__(self):
        super(Bug, self).__init__()
    def call(self, inputs):
        v, i = tf.nn.max_pool_with_argmax(inputs, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        return tf.cast(i, tf.float32)

b = Bug()
np.random.seed(0)
inp = np.random.uniform(0, 10, (2, 4, 4, 3))
tf_indices = b(inp)
print(inp[0,:,:,0])
print(tf_indices[0,:,:,0])
tf.saved_model.save(b, "bug")

To convert the model to onnx:

python -m tf2onnx.convert --saved-model bug --output bug.onnx --opset 13

To run it:

import onnxruntime as rt

session = rt.InferenceSession("bug.onnx")
input, output = session.get_inputs()[0], session.get_outputs()[0]
onnx_indices = session.run([output.name], {input.name: inp.astype(np.float32)})[0]
print(onnx_indices[0,:,:,0])

The tensorflow part of the code outputs:

[[5.48813504 5.44883183 4.37587211 3.83441519]
 [5.68044561 0.871293   7.78156751 7.99158564]
 [1.18274426 9.44668917 2.64555612 5.68433949]
 [6.12095723 6.81820299 6.97631196 6.7063787 ]]
tf.Tensor(
[[12. 21.]
 [27. 42.]], shape=(2, 2), dtype=float32)

While the onnxruntime one gives:

[[ 4.  7.]
 [ 9. 14.]]

Additional context
I also tried to convert such a simple model with keras-onnx, in which case I find that the onnx op returns both pooled values instead of pooled values and indices (there is an issue about this here).
Note that the indices are cast to float. I have let it because this caused an error when trying to convert with keras-onnx but it does not here so it can be ignored.

[TomWildenhain-Microsoft]

Hi @mpaillassa, thanks for the detailed report. You are correct, our implementation failed to convert the indices to the NHWC format.