Fix DecomonMaxPooling2D

src/decomon/layers/activations/activation.py

@@ -96,7 +96,10 @@ def forward_ibp_propagate(self, lower: Tensor, upper: Tensor) -> tuple[Tensor, T
         return self.layer.activation(lower), self.layer.activation(upper)
 
     def forward_affine_propagate(
-        self, input_affine_bounds: list[Tensor], input_constant_bounds: list[Tensor]
+        self,
+        input_affine_bounds: list[Tensor],
+        input_constant_bounds: list[Tensor],
+        perturbation_domain_inputs: list[Tensor],
     ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         if self.finetune and self.diagonal and len(input_affine_bounds) > 0:
             w_l_in, b_l_in, w_u_in, b_u_in = input_affine_bounds
@@ -151,7 +154,9 @@ def forward_affine_propagate(
             return (w_l_out, b_l_out, w_u_out, b_u_out)
         else:
             return super().forward_affine_propagate(
-                input_affine_bounds=input_affine_bounds, input_constant_bounds=input_constant_bounds
+                input_affine_bounds=input_affine_bounds,
+                input_constant_bounds=input_constant_bounds,
+                perturbation_domain_inputs=perturbation_domain_inputs,
             )
 
     def backward_affine_propagate(
@@ -287,8 +292,8 @@ def __init__(
         # so do the inputs/outputs format
         self.inputs_outputs_spec = self.decomon_activation.inputs_outputs_spec
 
-    def get_affine_representation(self) -> tuple[Tensor, Tensor]:
-        return self.decomon_activation.get_affine_representation()
+    def get_affine_representation(self, layer: Optional[Layer] = None) -> tuple[Tensor, Tensor]:
+        return self.decomon_activation.get_affine_representation(layer=layer)
 
     def get_affine_bounds(
         self,
@@ -299,10 +304,15 @@ def get_affine_bounds(
         return self.decomon_activation.get_affine_bounds(lower=lower, upper=upper)
 
     def forward_affine_propagate(
-        self, input_affine_bounds: list[Tensor], input_constant_bounds: list[Tensor]
+        self,
+        input_affine_bounds: list[Tensor],
+        input_constant_bounds: list[Tensor],
+        perturbation_domain_inputs: list[Tensor],
     ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         return self.decomon_activation.forward_affine_propagate(
-            input_affine_bounds=input_affine_bounds, input_constant_bounds=input_constant_bounds
+            input_affine_bounds=input_affine_bounds,
+            input_constant_bounds=input_constant_bounds,
+            perturbation_domain_inputs=perturbation_domain_inputs,
         )
 
     def backward_affine_propagate(

src/decomon/layers/backward/layer_backward.py

@@ -159,7 +159,10 @@ def forward_ibp_propagate(self, lower: Tensor, upper: Tensor) -> tuple[Tensor, T
             raise NotImplementedError()
 
     def forward_affine_propagate(
-        self, input_affine_bounds: list[Tensor], input_constant_bounds: list[Tensor]
+        self,
+        input_affine_bounds: list[Tensor],
+        input_constant_bounds: list[Tensor],
+        perturbation_domain_inputs: list[Tensor],
     ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         w_l_in, b_l_in, w_u_in, b_u_in = input_affine_bounds
         # reshape

src/decomon/layers/core/dense.py

@@ -1,7 +1,7 @@
-from typing import Any
+from typing import Any, Optional
 
 import keras.ops as K
-from keras.layers import Dense
+from keras.layers import Dense, Layer
 
 from decomon.layers.layer import DecomonLayer
 from decomon.types import Tensor
@@ -23,7 +23,7 @@ def __init__(
 
         super().__init__(*args, layer=layer, layer_pos=layer_pos, layer_neg=layer_neg, **kwargs)  # type: ignore
 
-    def get_affine_representation(self) -> tuple[Tensor, Tensor]:
+    def get_affine_representation(self, layer: Optional[Layer] = None) -> tuple[Tensor, Tensor]:
         w = self.layer.kernel
         b = self.layer.bias if self.layer.use_bias else K.zeros((self.layer.units,))
 

src/decomon/layers/custom/reduce/utils.py

@@ -137,9 +137,23 @@ def get_batch_multi_dot_repr_for_axis_reduce_weights(w: Tensor, axis: int, keepd
 
 
 def max_prime(inputs: Tensor, axis: int) -> Tensor:
-    indices = K.argmax(inputs, axis)
-    dim_i = inputs.shape[axis]
-    output = K.one_hot(indices, dim_i, axis=axis)
+    # preprocessing: need to overcome max nb of dimension for argmax (7 with tensorflow) => reshape
+    if axis < 0:
+        axis_ = len(inputs.shape) + axis
+    else:
+        axis_ = axis
+    oldshape = inputs.shape
+    newshape = (int(np.prod(inputs.shape[:axis_])), inputs.shape[axis_], int(np.prod(inputs.shape[axis_ + 1 :])))
+    inputs_reshaped = K.reshape(inputs, newshape=newshape)
+    axis_reshaped = 1
+
+    # max_prime: one-hot encoding of argmax
+    indices = K.argmax(inputs_reshaped, axis_reshaped)
+    dim_i = inputs_reshaped.shape[axis_reshaped]
+    output_reshaped = K.one_hot(indices, dim_i, axis=axis_reshaped)
+
+    # postprocessing: reshape back
+    output = K.reshape(output_reshaped, newshape=oldshape)
     return output
 
 

src/decomon/layers/layer.py

@@ -20,6 +20,7 @@
     get_affine_representation_wo_bias,
     get_bias,
 )
+from decomon.layers.utils.affine import get_affine_representation
 from decomon.perturbation_domain import BoxDomain, PerturbationDomain
 from decomon.types import Tensor
 from decomon.utils import fit_memory
@@ -110,6 +111,12 @@ class DecomonLayer(Wrapper):
     finetune_lower: bool = False
     "Flag telling that the layer can have its affine lower bound finetuned"
 
+    skip_forward_oracle: bool = False
+    """Flag to skip the forward_oracle in `call_forward()`.
+    In this case we keep passing perturbation_domain_inputs to `forward_affine_propagate`.
+    This can be useful with layers having linear parts like `MaxPooling2D`.
+    """
+
     def __init__(
         self,
         layer: Layer,
@@ -275,7 +282,7 @@ def get_affine_representation_upper(self) -> tuple[Tensor, Tensor]:
     def get_affine_representation_lower(self) -> tuple[Tensor, Tensor]:
         return self.get_affine_representation()
 
-    def get_affine_representation(self) -> tuple[Tensor, Tensor]:
+    def get_affine_representation(self, layer: Optional[Layer] = None) -> tuple[Tensor, Tensor]:
         """Get affine representation of the layer
 
         This computes the affine representation of the layer, when this is meaningful,
@@ -285,6 +292,7 @@ def get_affine_representation(self) -> tuple[Tensor, Tensor]:
         For non-linear layers, one should implement `get_affine_bounds()` instead.
 
         Args:
+            layer: linear component of a more complex layer. Default to `self.layer`.
 
         Returns:
             w, b: affine representation of the layer satisfying
@@ -315,13 +323,9 @@ def get_affine_representation(self) -> tuple[Tensor, Tensor]:
 
 
         """
-        if not self.linear:
-            raise RuntimeError("You should not call `get_affine_representation()` when `self.linear` is False.")
-        else:
-            if self.use_bias:
-                return get_affine_representation_with_bias(self.layer, diagonal=self.diagonal)
-            else:
-                return get_affine_representation_wo_bias(self.layer, diagonal=self.diagonal)
+        if layer is None:
+            layer = self.layer
+        return get_affine_representation(layer=layer, diagonal=self.diagonal, use_bias=self.use_bias)
 
     def get_affine_bounds(self, lower: Tensor, upper: Tensor, **kwargs: Any) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         """Get affine bounds on layer outputs from layer inputs
@@ -469,7 +473,7 @@ def _forward_affine_propagate_linear(
         """
         if len(input_affine_bounds) == 0:
             # special case: empty bounds <=> identity bounds
-            w, b = self.get_affine_representation()
+            w, b = self.get_affine_representation(layer=layer)
             return (w, b, w, b)
 
         w_l_in, b_l_in, w_u_in, b_u_in = input_affine_bounds
@@ -478,7 +482,7 @@ def _forward_affine_propagate_linear(
         is_from_diagonal = self.inputs_outputs_spec.is_diagonal_bounds(input_affine_bounds)
 
         if is_from_diagonal:
-            w_out, b_out = self.get_affine_representation()
+            w_out, b_out = self.get_affine_representation(layer=layer)
             layer_affine_bounds = [w_out, b_out] * 2
 
             from_linear_layer = (self.inputs_outputs_spec.is_wo_batch_bounds(input_affine_bounds), True)
@@ -562,7 +566,7 @@ def _forward_affine_propagate_linear(
 
                     return (w_l_out, b_l_out, w_u_out, b_u_out)
                 else:
-                    w, b = self.get_affine_representation()
+                    w, b = self.get_affine_representation(layer=layer)
                     layer_affine_bounds = [w, b, w, b]
                     from_linear_layer = (is_from_linear, self.linear)
                     diagonal = (
@@ -577,7 +581,10 @@ def _forward_affine_propagate_linear(
                     )
 
     def forward_affine_propagate(
-        self, input_affine_bounds: list[Tensor], input_constant_bounds: list[Tensor]
+        self,
+        input_affine_bounds: list[Tensor],
+        input_constant_bounds: list[Tensor],
+        perturbation_domain_inputs: list[Tensor],
     ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         """Propagate model affine bounds in forward direction.
 
@@ -589,6 +596,7 @@ def forward_affine_propagate(
                 affine bounds on underlying keras layer input w.r.t. model input
             input_constant_bounds: [l_c_in, u_c_in]
                 constant oracle bounds on underlying keras layer input (already deduced from affine ones if necessary)
+            perturbation_domain_inputs: perturbation domain input, wrapped in a list. Necessary only in self.skip_forward_oracle, else empty.
 
         Returns:
             w_l, b_l, w_u, b_u: affine bounds on underlying keras layer *output* w.r.t. model input
@@ -860,18 +868,27 @@ def call_forward(
         # Affine bounds propagation
         if self.affine:
             if not self.linear:
-                # get oracle input bounds (because input_bounds_to_propagate could be empty at this point)
-                input_constant_bounds = self.get_forward_oracle(
-                    input_affine_bounds=affine_bounds_to_propagate,
-                    input_constant_bounds=input_bounds_to_propagate,
-                    perturbation_domain_inputs=perturbation_domain_inputs,
-                )
+                if self.skip_forward_oracle:
+                    # we skip the oracle (should be done later during forward_affine_propagate)
+                    input_constant_bounds = input_bounds_to_propagate
+                    perturbation_domain_inputs_ = perturbation_domain_inputs
+                else:
+                    # get oracle input bounds (because input_bounds_to_propagate could be empty at this point)
+                    input_constant_bounds = self.get_forward_oracle(
+                        input_affine_bounds=affine_bounds_to_propagate,
+                        input_constant_bounds=input_bounds_to_propagate,
+                        perturbation_domain_inputs=perturbation_domain_inputs,
+                    )
+                    perturbation_domain_inputs_ = []
             else:
                 input_constant_bounds = []
+                perturbation_domain_inputs_ = []
             # forward propagation
             output_affine_bounds = list(
                 self.forward_affine_propagate(
-                    input_affine_bounds=affine_bounds_to_propagate, input_constant_bounds=input_constant_bounds
+                    input_affine_bounds=affine_bounds_to_propagate,
+                    input_constant_bounds=input_constant_bounds,
+                    perturbation_domain_inputs=perturbation_domain_inputs_,
                 )
             )
         else:

src/decomon/layers/merging/add.py

@@ -1,5 +1,7 @@
+from typing import Optional
+
 import keras.ops as K
-from keras.layers import Add
+from keras.layers import Add, Layer
 
 from decomon.layers.merging.base_merge import DecomonMerge
 from decomon.types import Tensor
@@ -11,7 +13,7 @@ class DecomonAdd(DecomonMerge):
     diagonal = True
     increasing = True
 
-    def get_affine_representation(self) -> tuple[list[Tensor], Tensor]:
+    def get_affine_representation(self, layer: Optional[Layer] = None) -> tuple[list[Tensor], Tensor]:
         w = [K.ones(input_i.shape[1:]) for input_i in self.keras_layer_input]
         b = K.zeros(self.layer.output.shape[1:])
 

src/decomon/layers/merging/average.py

@@ -1,5 +1,7 @@
+from typing import Optional
+
 import keras.ops as K
-from keras.layers import Average
+from keras.layers import Average, Layer
 
 from decomon.layers.merging.base_merge import DecomonMerge
 from decomon.types import Tensor
@@ -11,7 +13,7 @@ class DecomonAverage(DecomonMerge):
     diagonal = True
     increasing = True
 
-    def get_affine_representation(self) -> tuple[list[Tensor], Tensor]:
+    def get_affine_representation(self, layer: Optional[Layer] = None) -> tuple[list[Tensor], Tensor]:
         coeff = 1.0 / len(self.keras_layer_input)
         w = [coeff * K.ones(input_i.shape[1:]) for input_i in self.keras_layer_input]
         b = K.zeros(self.layer.output.shape[1:])

src/decomon/layers/merging/base_merge.py

@@ -2,6 +2,7 @@
 
 import keras
 import keras.ops as K
+from keras import Layer
 
 from decomon.keras_utils import add_tensors, batch_multid_dot
 from decomon.layers.fuse import combine_affine_bounds
@@ -37,7 +38,7 @@ def layer_input_shape_wo_batchsize(self) -> list[list[int]]:  # type: ignore
         else:
             return [list(e.shape[1:]) for e in self.layer.input]
 
-    def get_affine_representation(self) -> tuple[list[Tensor], Tensor]:  # type: ignore
+    def get_affine_representation(self, layer: Optional[Layer] = None) -> tuple[list[Tensor], Tensor]:  # type: ignore
         """Get affine representation of the layer
 
         This computes the affine representation of the layer, when this is meaningful,
@@ -199,7 +200,10 @@ def forward_ibp_propagate(self, lower: list[Tensor], upper: list[Tensor]) -> tup
             )
 
     def forward_affine_propagate(
-        self, input_affine_bounds: list[list[Tensor]], input_constant_bounds: list[list[Tensor]]
+        self,
+        input_affine_bounds: list[list[Tensor]],
+        input_constant_bounds: list[list[Tensor]],
+        perturbation_domain_inputs: list[Tensor],
     ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         """Propagate model affine bounds in forward direction.
 
@@ -211,6 +215,7 @@ def forward_affine_propagate(
                 affine bounds on underlying keras layer i-th input w.r.t. model input
             input_constant_bounds[i]: [l_c_in[i], u_c_in[i]]
                 constant oracle bounds on underlying keras layer i-th input (already deduced from affine ones if necessary)
+            perturbation_domain_inputs: perturbation domain input, wrapped in a list. Necessary only in self.skip_forward_oracle, else empty.
 
         Returns:
             w_l_new, b_l_new, w_u_new, b_u_new: affine bounds on underlying keras layer *output* w.r.t. model input

src/decomon/layers/merging/subtract.py

@@ -1,5 +1,7 @@
+from typing import Optional
+
 import keras.ops as K
-from keras.layers import Subtract
+from keras.layers import Layer, Subtract
 
 from decomon.layers.merging.base_merge import DecomonMerge
 from decomon.types import Tensor
@@ -10,7 +12,7 @@ class DecomonSubtract(DecomonMerge):
     linear = True
     diagonal = True
 
-    def get_affine_representation(self) -> tuple[list[Tensor], Tensor]:
+    def get_affine_representation(self, layer: Optional[Layer] = None) -> tuple[list[Tensor], Tensor]:
         w = [K.ones(input_i.shape[1:]) for input_i in self.keras_layer_input]
         w[1] *= -1
         b = K.zeros(self.layer.output.shape[1:])

src/decomon/layers/normalization/spectral_normalization.py

@@ -47,17 +47,22 @@ def __init__(
         self.sub_layer: Layer = self.layer.layer
         self.decomon_layer = self.layer
 
-    def get_affine_representation(self) -> tuple[Tensor, Tensor]:
-        return self.decomon_layer.get_affine_representation()
+    def get_affine_representation(self, layer: Optional[Layer] = None) -> tuple[Tensor, Tensor]:
+        return self.decomon_layer.get_affine_representation(layer=layer)
 
     def forward_ibp_propagate(self, lower: Tensor, upper: Tensor) -> tuple[Tensor, Tensor]:
         return self.decomon_layer.forward_ibp_propagate(lower, upper)
 
     def forward_affine_propagate(
-        self, input_affine_bounds: list[Tensor], input_constant_bounds: list[Tensor]
+        self,
+        input_affine_bounds: list[Tensor],
+        input_constant_bounds: list[Tensor],
+        perturbation_domain_inputs: list[Tensor],
     ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         return self.decomon_layer.forward_affine_propagate(
-            input_affine_bounds=input_affine_bounds, input_constant_bounds=input_constant_bounds
+            input_affine_bounds=input_affine_bounds,
+            input_constant_bounds=input_constant_bounds,
+            perturbation_domain_inputs=perturbation_domain_inputs,
         )
 
     def backward_affine_propagate(