fix some review comments

Author: stes

cebra/attribution/jacobian_attribution.py

@@ -52,14 +52,30 @@ def get_attribution_map(
     convert_to_numpy: bool = True,
     aggregate: Literal["mean", "sum", "max"] = "mean",
     transform: Literal["none", "abs"] = "none",
-    hybrid_solver=False,
+    hybrid_solver: bool = False,
 ):
-    """Estimate attribution maps.
+    """Estimate attribution maps using the Jacobian pseudo-inverse.
+
     The function estimates Jacobian matrices for each point in the model,
-    computes the pseudo-inverse (for every sample), applies the `transform`
-    function point-wise, and then aggregates with the `aggregate` function
-    over the sample dimension.
-    The result is a `(num_inputs, num_features)` attribution map.
+    computes the pseudo-inverse (for every sample) and then aggregates
+    the resulting matrices to compute an attribution map.
+
+    Args:
+        model: The neural network model for which to compute attributions.
+        input_data: Input tensor or numpy array to compute attributions for.
+        double_precision: If ``True``, use double precision for computation.
+        convert_to_numpy: If ``True``, convert the output to numpy arrays.
+        aggregate: Method to aggregate attribution values across samples.
+            Options are ``"mean"``, ``"sum"``, or ``"max"``.
+        transform: Transformation to apply to attribution values.
+            Options are ``"none"`` or ``"abs"``.
+        hybrid_solver: If ``True``, handle multi-objective models differently.
+
+    Returns:
+        A tuple containing:
+            - jf: The Jacobian matrix of shape (num_samples, output_dim, input_dim)
+            - jhatg: The pseudo-inverse of the Jacobian matrix
+        The result is effectively a ``(num_inputs, num_features)`` attribution map.
     """
     assert aggregate in ["mean", "sum", "max"]
 

cebra/models/multiobjective.py

@@ -283,13 +283,14 @@ def __init__(self,
 
         if max_slice_dim != self.num_output:
             raise ValueError(
-                f"The dimension of output {self.num_output} is different than the highest dimension of slices {max_slice_dim}."
-                f"They need to have the same dimension.")
+                f"The dimension of output {self.num_output} is different than the highest dimension of the slices ({max_slice_dim})."
+                f"The output dimension and slice dimension need to have the same dimension."
+            )
 
         check_slices_for_gaps(self.feature_ranges)
 
         if check_overlapping_feature_ranges(self.feature_ranges):
-            print("Computing renormalize ranges...")
+            print("Computing renormalized ranges...")
             self.renormalize_ranges = compute_renormalize_ranges(
                 self.feature_ranges, sort=True)
             print("New ranges:", self.renormalize_ranges)
@@ -327,9 +328,12 @@ def forward(self, inputs):
 
         if self.renormalize:
             if hasattr(self, "renormalize_ranges"):
-                #TODO: does the order of the renormalize ranges matter??
-                # I think it does, imagine that the renormalize ranges are (5, 10), (0, 5), then
-                # when we do torch.cat() output will be wrong --> Renormalize ranges need to be ordered.
+                if not all(self.renormalize_ranges[i].start <=
+                           self.renormalize_ranges[i + 1].start
+                           for i in range(len(self.renormalize_ranges) - 1)):
+                    raise ValueError(
+                        "The renormalize_ranges must be sorted by start index.")
+
                 output = [
                     self._norm(output[:, slice_features])
                     for slice_features in self.renormalize_ranges