Add adaptive pooling (1D, 2D, 3D) support across JAX, NumPy, TensorFlow, and PyTorch backends (#21820)

* Add AdaptiveAveragePooling2D and AdaptiveMaxPooling2D layers

* Add adaptive pooling (adaptive_avg_pool and adaptive_max_pool) for JAX, NumPy, PyTorch, and TensorFlow backends

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Fix adaptive pooling implementation

* Refactor adaptive pooling with shared utils and base classes

* Update adaptive pooling implementation per review feedback

* Update adaptive pooling implementation per review feedback

* Fix config imports and regenerate API.

* Fix exports.

* Fix tests, in particular with Torch on GPU.

---------

Co-authored-by: Fabien Hertschuh <1091026+hertschuh@users.noreply.github.com>

Author: MalyalaKarthik66

keras/api/_tf_keras/keras/layers/__init__.py

@@ -113,6 +113,24 @@
 from keras.src.layers.normalization.unit_normalization import (
     UnitNormalization as UnitNormalization,
 )
+from keras.src.layers.pooling.adaptive_average_pooling1d import (
+    AdaptiveAveragePooling1D as AdaptiveAveragePooling1D,
+)
+from keras.src.layers.pooling.adaptive_average_pooling2d import (
+    AdaptiveAveragePooling2D as AdaptiveAveragePooling2D,
+)
+from keras.src.layers.pooling.adaptive_average_pooling3d import (
+    AdaptiveAveragePooling3D as AdaptiveAveragePooling3D,
+)
+from keras.src.layers.pooling.adaptive_max_pooling1d import (
+    AdaptiveMaxPooling1D as AdaptiveMaxPooling1D,
+)
+from keras.src.layers.pooling.adaptive_max_pooling2d import (
+    AdaptiveMaxPooling2D as AdaptiveMaxPooling2D,
+)
+from keras.src.layers.pooling.adaptive_max_pooling3d import (
+    AdaptiveMaxPooling3D as AdaptiveMaxPooling3D,
+)
 from keras.src.layers.pooling.average_pooling1d import (
     AveragePooling1D as AveragePooling1D,
 )

keras/api/_tf_keras/keras/ops/__init__.py

@@ -64,6 +64,8 @@
 from keras.src.ops.math import top_k as top_k
 from keras.src.ops.math import view_as_complex as view_as_complex
 from keras.src.ops.math import view_as_real as view_as_real
+from keras.src.ops.nn import adaptive_average_pool as adaptive_average_pool
+from keras.src.ops.nn import adaptive_max_pool as adaptive_max_pool
 from keras.src.ops.nn import average_pool as average_pool
 from keras.src.ops.nn import batch_normalization as batch_normalization
 from keras.src.ops.nn import binary_crossentropy as binary_crossentropy

keras/api/_tf_keras/keras/ops/nn/__init__.py

@@ -4,6 +4,8 @@
 since your modifications would be overwritten.
 """
 
+from keras.src.ops.nn import adaptive_average_pool as adaptive_average_pool
+from keras.src.ops.nn import adaptive_max_pool as adaptive_max_pool
 from keras.src.ops.nn import average_pool as average_pool
 from keras.src.ops.nn import batch_normalization as batch_normalization
 from keras.src.ops.nn import binary_crossentropy as binary_crossentropy

keras/api/layers/__init__.py

@@ -113,6 +113,24 @@
 from keras.src.layers.normalization.unit_normalization import (
     UnitNormalization as UnitNormalization,
 )
+from keras.src.layers.pooling.adaptive_average_pooling1d import (
+    AdaptiveAveragePooling1D as AdaptiveAveragePooling1D,
+)
+from keras.src.layers.pooling.adaptive_average_pooling2d import (
+    AdaptiveAveragePooling2D as AdaptiveAveragePooling2D,
+)
+from keras.src.layers.pooling.adaptive_average_pooling3d import (
+    AdaptiveAveragePooling3D as AdaptiveAveragePooling3D,
+)
+from keras.src.layers.pooling.adaptive_max_pooling1d import (
+    AdaptiveMaxPooling1D as AdaptiveMaxPooling1D,
+)
+from keras.src.layers.pooling.adaptive_max_pooling2d import (
+    AdaptiveMaxPooling2D as AdaptiveMaxPooling2D,
+)
+from keras.src.layers.pooling.adaptive_max_pooling3d import (
+    AdaptiveMaxPooling3D as AdaptiveMaxPooling3D,
+)
 from keras.src.layers.pooling.average_pooling1d import (
     AveragePooling1D as AveragePooling1D,
 )

keras/api/ops/__init__.py

@@ -64,6 +64,8 @@
 from keras.src.ops.math import top_k as top_k
 from keras.src.ops.math import view_as_complex as view_as_complex
 from keras.src.ops.math import view_as_real as view_as_real
+from keras.src.ops.nn import adaptive_average_pool as adaptive_average_pool
+from keras.src.ops.nn import adaptive_max_pool as adaptive_max_pool
 from keras.src.ops.nn import average_pool as average_pool
 from keras.src.ops.nn import batch_normalization as batch_normalization
 from keras.src.ops.nn import binary_crossentropy as binary_crossentropy

keras/api/ops/nn/__init__.py

@@ -4,6 +4,8 @@
 since your modifications would be overwritten.
 """
 
+from keras.src.ops.nn import adaptive_average_pool as adaptive_average_pool
+from keras.src.ops.nn import adaptive_max_pool as adaptive_max_pool
 from keras.src.ops.nn import average_pool as average_pool
 from keras.src.ops.nn import batch_normalization as batch_normalization
 from keras.src.ops.nn import binary_crossentropy as binary_crossentropy

keras/src/backend/common/backend_utils.py

@@ -1,4 +1,5 @@
 import functools
+import math
 import operator
 import re
 import warnings
@@ -539,3 +540,10 @@ def slice_along_axis(x, start=0, stop=None, step=1, axis=0):
             -1 - axis
         )
     return x[tuple(slices)]
+
+
+def compute_adaptive_pooling_window_sizes(input_dim, output_dim):
+    """Compute small and big window sizes for adaptive pooling."""
+    small = math.ceil(input_dim / output_dim)
+    big = small + 1
+    return small, big