tf.compat.v1.nn.max_pool_with_argmax