tf.nn.max_pool2d: No support for tf.Tensor input for ksize

[luuzk]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): no
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10.0.19041
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: -
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.3.0
• Python version: 3.7.6
• Bazel version (if compiling from source): -
• GCC/Compiler version (if compiling from source): -
• CUDA/cuDNN version: 11.0/7.6.5
• GPU model and memory: RTX 2070 8GB

Describe the current behavior
The ksize arguement of tf.nn.max_pool2d supports only native Python int but not tf.Tensor. This is not a problem when using eager mode, but very disadvantageous when calculating the ksize as a Tensor in TF Graph mode.

TypeError: Expected int for argument 'ksize' not <tf.Tensor 'random_uniform:0' shape=() dtype=int32>.

Describe the expected behavior
Tensor support for tf.nn.max_pool2d function parameters, espacially ksize.

Standalone code to reproduce the issue

@tf.function
def issue():
    inp = tf.convert_to_tensor(np.arange(75).reshape((1, 5, 5, 3)))
    ksize = tf.random.uniform([], 1, 3, dtype=tf.int32)
    return tf.nn.max_pool2d(inp, [ksize], 1, "SAME")

if __name__ == "__main__":
    issue()

Other info / logs
traceback.txt

Thank you in advance!

[bhack]

Check https://stackoverflow.com/questions/43574076/tensorflow-maxpool-with-dynamic-ksize
And #4746

[luuzk]

Thank you @bhack! There seems to be a solution for this issue provided by PR #11875. Use gen_nn_ops.max_pool_v2 instead of gen_nn_ops.max_pool, which is used by default when calling tf.nn.max_pool2d. Change the code above to:

from tensorflow.python.ops import gen_nn_ops
return gen_nn_ops.max_pool_v2(inp, [1, ksize, ksize, 1], [1, 1, 1, 1], "SAME")

PR #14983 tries to integrate the utilization of gen_nn_ops.max_pool_v2 as standard behavior but was reverted. @yongtang @rmlarsen, are you interested to push PR #14983 into tensorflow:master to completely solve this issue, #1967 and, #4746?

[sachinprasadhs]

Was able to reproduce your issue in Tensorflow 2.5 version, please find the gist here. Thanks!

[cl886699]

from tensorflow.python.ops import gen_nn_ops
return gen_nn_ops.max_pool_v2(inp, [1, ksize, ksize, 1], [1, 1, 1, 1], "SAME")

this is not working for me

[chunduriv]

Was able to reproduce your issue in Tensorflow 2.11 version, please find the gist here. Thanks!

[bhack]

/cc @rmlarsen @yongtang