Torch/Paddle adaptive_max_pool2d (ivy-llc#21247)

ivy/data_classes/array/experimental/layers.py

@@ -832,6 +832,33 @@ def adaptive_avg_pool2d(
  output_size,
  )
 
+ def adaptive_max_pool2d(
+ self: ivy.Array,
+ output_size: Union[Sequence[int], int],
+ ) -> ivy.Array:
+ """
+ Apply a 2D adaptive maximum pooling over an input signal composed of several
+ input planes.
+
+ Parameters
+ ----------
+ self
+ Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N
+ is the batch dimension, C is the feature dimension, and H_in and W_in are
+ the 2 spatial dimensions.
+ output_size
+ Spatial output size.
+
+ Returns
+ -------
+ The result of the pooling operation. Will have shape (N, C, S_0, S_1) or
+ (C, S_0, S_1), where S = `output_size`
+ """
+ return ivy.adaptive_max_pool2d(
+ self._data,
+ output_size,
+ )
+
  def reduce_window(
  self: ivy.Array,
  init_value: Union[int, float],

ivy/data_classes/container/experimental/layers.py

@@ -1910,6 +1910,78 @@ def adaptive_avg_pool2d(
  map_sequences=map_sequences,
  )
 
+ @staticmethod
+ def static_adaptive_max_pool2d(
+ input: Union[ivy.Array, ivy.NativeArray, ivy.Container],
+ output_size: Union[Sequence[int], int, ivy.Container],
+ *,
+ key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None,
+ to_apply: Union[bool, ivy.Container] = True,
+ prune_unapplied: Union[bool, ivy.Container] = False,
+ map_sequences: Union[bool, ivy.Container] = False,
+ ) -> ivy.Container:
+ """
+ ivy.Container static method variant of ivy.adaptive_max_pool2d. This method
+ simply wraps the function, and so the docstring for ivy.adaptive_max_pool2d also
+ applies to this method with minimal changes.
+
+ Parameters
+ ----------
+ input
+ Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N
+ is the batch dimension, C is the feature dimension, and H_in and W_in are
+ the 2 spatial dimensions.
+ output_size
+ Spatial output size.
+
+ Returns
+ -------
+ The result of the pooling operation. Will have shape (N, C, S_0, S_1) or
+ (C, S_0, S_1), where S = `output_size`
+ """
+ return ContainerBase.cont_multi_map_in_function(
+ "adaptive_max_pool2d",
+ input,
+ output_size,
+ key_chains=key_chains,
+ to_apply=to_apply,
+ prune_unapplied=prune_unapplied,
+ map_sequences=map_sequences,
+ )
+
+ def adaptive_max_pool2d(
+ self: ivy.Container,
+ output_size: Union[int, ivy.Container],
+ *,
+ key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None,
+ to_apply: Union[bool, ivy.Container] = True,
+ prune_unapplied: Union[bool, ivy.Container] = False,
+ map_sequences: Union[bool, ivy.Container] = False,
+ ) -> ivy.Container:
+ """
+ Apply a 2D adaptive maximum pooling over an input signal composed of several
+ input planes.
+
+ Parameters
+ ----------
+ self
+ Input container.
+ output_size
+ Spatial output size.
+
+ Returns
+ -------
+ The result of the pooling operation.
+ """
+ return self.static_adaptive_max_pool2d(
+ self,
+ output_size,
+ key_chains=key_chains,
+ to_apply=to_apply,
+ prune_unapplied=prune_unapplied,
+ map_sequences=map_sequences,
+ )
+
  @staticmethod
  def static_ifftn(
  x: ivy.Container,

ivy/functional/backends/paddle/experimental/layers.py

@@ -381,6 +381,15 @@ def interpolate(
  raise IvyNotImplementedException()
 
 
+def adaptive_max_pool2d(
+ input: paddle.Tensor, output_size: Union[Sequence[int], int]
+) -> paddle.Tensor:
+ squeeze = input.ndim == 3
+ x = paddle.unsqueeze(input, axis=0) if squeeze else input
+ ret = paddle.nn.functional.adaptive_max_pool2d(x, output_size)
+ return paddle.squeeze(ret, axis=0) if squeeze else ret
+
+
 def ifftn(
  x: paddle.Tensor,
  s: Optional[Union[int, Tuple[int]]] = None,

ivy/functional/backends/torch/experimental/layers.py

@@ -865,6 +865,13 @@ def interpolate(
 ]
 
 
+@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version)
+def adaptive_max_pool2d(
+ input: torch.Tensor, output_size: Union[Sequence[int], int]
+) -> torch.Tensor:
+ return torch.nn.functional.adaptive_max_pool2d(input, output_size)
+
+
 @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version)
 def adaptive_avg_pool1d(input, output_size):
  return torch.nn.functional.adaptive_avg_pool1d(input, output_size)

ivy/functional/frontends/paddle/nn/functional/pooling.py

@@ -83,6 +83,12 @@ def adaptive_avg_pool2d(x, output_size, data_format="NCHW", name=None):
  return ivy.adaptive_avg_pool2d(x, output_size)
 
 
+@to_ivy_arrays_and_back
+@with_supported_dtypes({"2.5.1 and below": ("float32", "float64")}, "paddle")
+def adaptive_max_pool2d(x, output_size, return_mask=None, name=None):
+ return ivy.adaptive_max_pool2d(x, output_size)
+
+
 @to_ivy_arrays_and_back
 @with_supported_dtypes({"2.5.1 and below": ("float32", "float64")}, "paddle")
 def max_unpool1d(

ivy/functional/frontends/torch/nn/functional/pooling_functions.py

@@ -159,6 +159,16 @@ def max_pool2d(
  return ret
 
 
+@to_ivy_arrays_and_back
+def adaptive_max_pool2d(
+ input,
+ output_size,
+ return_indices=False,
+):
+ # ToDo: Add return_indices once superset is implemented
+ return ivy.adaptive_max_pool2d(input, output_size)
+
+
 @with_unsupported_dtypes(
  {
  "2.0.1 and below": (

ivy/functional/ivy/experimental/layers.py

@@ -2081,19 +2081,111 @@ def _expand_to_dim(x, dim):
  return x
 
 
-def _mask(vals, length, range_max, dim):
+def _mask(vals, length, range_max, dim, mask_value=0.0):
  if isinstance(length, int):
  return vals, length
  else:
  assert dim < 0
  mask = ivy.greater_equal(range_max, ivy.expand_dims(length, axis=-1))
  if dim == -2:
  mask = _expand_to_dim(mask, 4)
- vals = ivy.where(mask, 0.0, vals)
+ vals = ivy.where(mask, ivy.array(mask_value, device=vals.device), vals)
  length = _expand_to_dim(length, -dim)
  return vals, length
 
 
+@handle_nestable
+@inputs_to_ivy_arrays
+def adaptive_max_pool2d(
+ input: Union[ivy.Array, ivy.NativeArray],
+ output_size: Union[Sequence[int], int],
+):
+ """
+ Apply a 2D adaptive maximum pooling over an input signal composed of several input
+ planes.
+
+ Parameters
+ ----------
+ input
+ Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is
+ the batch dimension, C is the feature dimension, and H_in and W_in are the 2
+ spatial dimensions.
+ output_size
+ Spatial output size.
+
+ Returns
+ -------
+ The result of the pooling operation. Will have shape (N, C, S_0, S_1) or
+ (C, S_0, S_1), where S = `output_size`
+ """
+ squeeze = False
+ if input.ndim == 3:
+ input = ivy.expand_dims(input, axis=0)
+ squeeze = True
+ elif input.ndim != 4:
+ raise ivy.utils.exceptions.IvyException(
+ f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.",
+ )
+
+ if isinstance(output_size, int):
+ output_size = (output_size, output_size)
+
+ if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)):
+ stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size))
+ kernel_size = stride # Mathematically identical to the previous expression
+ pooled_output = ivy.max_pool2d(
+ input, kernel_size, stride, "VALID", data_format="NCHW"
+ )
+ if squeeze:
+ return ivy.squeeze(pooled_output, axis=0)
+ return pooled_output
+
+ idxh, length_h, range_max_h, adaptive_h = _compute_idx(
+ input.shape[-2], output_size[-2], input.device
+ )
+ idxw, length_w, range_max_w, adaptive_w = _compute_idx(
+ input.shape[-1], output_size[-1], input.device
+ )
+
+ # to numpy and back in order to bypass a slicing error in tensorflow
+ vals = ivy.array(
+ input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw], device=input.device
+ )
+
+ if not adaptive_h and not adaptive_w:
+ ret = ivy.max(vals, axis=(-3, -1))
+ ret = ivy.squeeze(ret, axis=0) if squeeze else ret
+ return ret
+
+ vals, length_h = _mask(
+ vals, length_h, range_max_h, dim=-2, mask_value=float("-inf")
+ )
+ vals, length_w = _mask(
+ vals, length_w, range_max_w, dim=-1, mask_value=float("-inf")
+ )
+
+ ret = None
+ for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])):
+ if ret is None:
+ ret = vals[..., i, :, j]
+ else:
+ ret = ivy.maximum(ret, vals[..., i, :, j])
+ pooled_output = ret.astype(vals.dtype)
+
+ pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output
+ return pooled_output
+
+
+adaptive_max_pool2d.mixed_backend_wrappers = {
+ "to_add": (
+ "inputs_to_native_arrays",
+ "outputs_to_ivy_arrays",
+ "handle_device_shifting",
+ ),
+ "to_skip": ("inputs_to_ivy_arrays",),
+}
+
+
 @handle_nestable
 @inputs_to_ivy_arrays
 def adaptive_avg_pool1d(
@@ -2214,10 +2306,7 @@ def adaptive_avg_pool2d(
 
  if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)):
  stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size))
- kernel_size = tuple(
- i_s - (o_s - 1) * st
- for i_s, o_s, st in zip(input.shape[-2:], output_size, stride)
- )
+ kernel_size = stride # Mathematically identical to the previous expression
  pooled_output = ivy.avg_pool2d(
  input, kernel_size, stride, "VALID", data_format="NCHW"
  )

ivy_tests/test_ivy/test_frontends/test_paddle/test_nn/test_functional/test_pooling.py

@@ -190,6 +190,50 @@ def test_paddle_adaptive_avg_pool2d(
  )
 
 
+# adaptive_max_pool2d
+@handle_frontend_test(
+ fn_tree="paddle.nn.functional.adaptive_max_pool2d",
+ dtype_and_x=helpers.dtype_and_values(
+ available_dtypes=helpers.get_dtypes("float"),
+ min_num_dims=4,
+ max_num_dims=4,
+ min_dim_size=1,
+ # Setting max and min value because this operation in paddle is not
+ # numerically stable
+ max_value=100,
+ min_value=-100,
+ ),
+ output_size=st.one_of(
+ st.tuples(
+ helpers.ints(min_value=1, max_value=5),
+ helpers.ints(min_value=1, max_value=5),
+ ),
+ helpers.ints(min_value=1, max_value=5),
+ ),
+)
+def test_paddle_adaptive_max_pool2d(
+ *,
+ dtype_and_x,
+ output_size,
+ test_flags,
+ frontend,
+ on_device,
+ backend_fw,
+ fn_tree,
+):
+ input_dtype, x = dtype_and_x
+ helpers.test_frontend_function(
+ input_dtypes=input_dtype,
+ backend_to_test=backend_fw,
+ frontend=frontend,
+ test_flags=test_flags,
+ on_device=on_device,
+ fn_tree=fn_tree,
+ x=x[0],
+ output_size=output_size,
+ )
+
+
 # max_unpool1d
 @handle_frontend_test(
  fn_tree="paddle.nn.functional.max_unpool1d",