Fix decomon custom layers

src/decomon/keras_utils.py

@@ -20,7 +20,7 @@ def batch_multid_dot(
 ) -> Tensor:
     """Dot product of tensors by batch, along multiple axes
 
-    Hypothesis: we sum over last axes of x and first axes (skipping the batch one) of x.
+    Hypothesis: we sum over last axes of x and first axes (skipping the batch one) of y.
 
     The 1-dimensional equivalent would be `batch_dot(x,y, axes=(-1, 1))`
     or `keras.layers.Dot(axes=(-1, 1))(x,y)`

src/decomon/layers/custom/__init__.py

@@ -1,4 +1,2 @@
-from .max import DecomonMax
-from .min import DecomonMin
-from .mul import DecomonMulConstant
-from .utils import get_affine_lower_bound_max, get_affine_upper_bound_max
+from .constant import DecomonMulConstant
+from .reduce import DecomonMax, DecomonMin

src/decomon/layers/custom/constant/__init__.py

@@ -0,0 +1 @@
+from .constant import DecomonMulConstant

src/decomon/layers/custom/constant/constant.py

@@ -0,0 +1,10 @@
+from keras_custom.layers import MulConstant
+
+from decomon.layers.layer import DecomonLayer
+
+
+class DecomonMulConstant(DecomonLayer):
+    layer: MulConstant
+    linear = True
+    diagonal = True
+    use_bias = False

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

@@ -0,0 +1,6 @@
+from .max import DecomonMax
+from .min import DecomonMin
+from .utils import (
+    get_affine_lower_bound_max_before_reduction,
+    get_affine_upper_bound_max_before_reduction,
+)

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

@@ -0,0 +1,37 @@
+# define non native class Max
+# Decomon Custom for Max(axis...)
+from typing import Any
+
+from keras_custom.layers import Max
+
+from decomon.layers.layer import DecomonLayer
+from decomon.types import Tensor
+
+from .utils import (
+    get_affine_lower_bound_max_before_reduction,
+    get_affine_upper_bound_max_before_reduction,
+    get_batch_multi_dot_repr_for_axis_reduce_weights,
+)
+
+
+class DecomonMax(DecomonLayer):
+    layer: Max
+    linear = False
+    increasing = True
+
+    def get_affine_bounds(self, lower: Tensor, upper: Tensor, **kwargs: Any) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        return get_affine_bounds_max(lower=lower, upper=upper, axis=self.layer.axis, keepdims=self.layer.keepdims)
+
+
+def get_affine_bounds_max(
+    lower: Tensor, upper: Tensor, axis: int, keepdims: bool
+) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+    w_l, b_l = get_affine_lower_bound_max_before_reduction(lower, upper, axis=axis, keepdims=keepdims)
+    w_u, b_u = get_affine_upper_bound_max_before_reduction(lower, upper, axis=axis, keepdims=keepdims)
+
+    # for now the lower bound is given by sum(w_l*x, axis=axis) + b,
+    # so we need another transformation to get the final w_l
+    w_l = get_batch_multi_dot_repr_for_axis_reduce_weights(w_l, axis=axis, keepdims=keepdims)
+    w_u = get_batch_multi_dot_repr_for_axis_reduce_weights(w_u, axis=axis, keepdims=keepdims)
+
+    return (w_l, b_l, w_u, b_u)