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

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