congrads 1.0.7__py3-none-any.whl → 1.1.0__py3-none-any.whl

congrads/constraints.py

@@ -1,65 +1,55 @@
-"""
-This module provides a set of constraint classes for guiding neural network 
-training by enforcing specific conditions on the network's outputs.
-
-The constraints in this module include:
-
-- `Constraint`: The base class for all constraint types, defining the 
-  interface and core behavior.
-- `ImplicationConstraint`: A constraint that enforces one condition only if 
-  another condition is met, useful for modeling implications between network 
-  outputs.
-- `ScalarConstraint`: A constraint that enforces scalar-based comparisons on 
-  a network's output.
-- `BinaryConstraint`: A constraint that enforces a binary comparison between
-  two neurons in the network, using a comparison function (e.g., less than,
-  greater than).
-- `SumConstraint`: A constraint that enforces that the sum of certain neurons'
-  outputs equals a specified value, which can be used to control total output.
-- `PythagoreanConstraint`: A constraint that enforces the Pythagorean theorem
-  on a set of neurons, ensuring that the square of one neuron's output is equal
-  to the sum of the squares of other outputs.
-
-These constraints can be used to steer the learning process by applying 
-conditions such as logical implications or numerical bounds.
+"""Module providing constraint classes for guiding neural network training.
+
+This module defines constraints that enforce specific conditions on network outputs
+to steer learning. Available constraint types include:
+
+- `Constraint`: Base class for all constraint types, defining the interface and core
+  behavior.
+- `ImplicationConstraint`: Enforces one condition only if another condition is met,
+  useful for modeling implications between outputs.
+- `ScalarConstraint`: Enforces scalar-based comparisons on a network's output.
+- `BinaryConstraint`: Enforces a binary comparison between two tags using a
+  comparison function (e.g., less than, greater than).
+- `SumConstraint`: Ensures the sum of selected tags' outputs equals a specified
+  value, controlling total output.
+
+These constraints can steer the learning process by applying logical implications
+or numerical bounds.
 
 Usage:
     1. Define a custom constraint class by inheriting from `Constraint`.
-    2. Apply the constraint to your neural network during training to 
-       enforce desired output behaviors.
-    3. Use the helper classes like `IdentityTransformation` for handling 
-       transformations and comparisons in constraints.
+    2. Apply the constraint to your neural network during training.
+    3. Use helper classes like `IdentityTransformation` for transformations and
+       comparisons in constraints.
 
-Dependencies:
-    - PyTorch (`torch`)
 """
 
 import random
 import string
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from numbers import Number
-from typing import Callable, Dict, Union
+from typing import Literal
 
 from torch import (
     Tensor,
-    count_nonzero,
+    argsort,
+    eq,
     ge,
     gt,
-    isclose,
     le,
+    logical_and,
     logical_not,
     logical_or,
     lt,
-    numel,
     ones,
     ones_like,
     reshape,
     sign,
-    sqrt,
-    square,
     stack,
     tensor,
+    unique,
     zeros_like,
 )
 from torch.nn.functional import normalize
@@ -70,8 +60,7 @@ from .utils import validate_comparator_pytorch, validate_iterable, validate_type
 
 
 class Constraint(ABC):
-    """
-    Abstract base class for defining constraints applied to neural networks.
+    """Abstract base class for defining constraints applied to neural networks.
 
     A `Constraint` specifies conditions that the neural network outputs
     should satisfy. It supports monitoring constraint satisfaction
@@ -79,18 +68,18 @@ class Constraint(ABC):
     must implement the `check_constraint` and `calculate_direction` methods.
 
     Args:
-        neurons (set[str]): Names of the neurons this constraint applies to.
+        tags (set[str]): Tags referencing parts of the network where this constraint applies to.
         name (str, optional): A unique name for the constraint. If not provided,
             a name is generated based on the class name and a random suffix.
-        monitor_only (bool, optional): If True, only monitor the constraint
-            without adjusting the loss. Defaults to False.
+        enforce (bool, optional): If False, only monitor the constraint
+            without adjusting the loss. Defaults to True.
         rescale_factor (Number, optional): Factor to scale the
             constraint-adjusted loss. Defaults to 1.5. Should be greater
             than 1 to give weight to the constraint.
 
     Raises:
         TypeError: If a provided attribute has an incompatible type.
-        ValueError: If any neuron in `neurons` is not
+        ValueError: If any tag in `tags` is not
             defined in the `descriptor`.
 
     Note:
@@ -104,29 +93,44 @@ class Constraint(ABC):
     device = None
 
     def __init__(
-        self,
-        neurons: set[str],
-        name: str = None,
-        monitor_only: bool = False,
-        rescale_factor: Number = 1.5,
+        self, tags: set[str], name: str = None, enforce: bool = True, rescale_factor: Number = 1.5
     ) -> None:
-        """
-        Initializes a new Constraint instance.
-        """
+        """Initializes a new Constraint instance.
 
+        Args:
+            tags (set[str]): Tags referencing parts of the network where this constraint applies to.
+            name (str, optional): A unique name for the constraint. If not
+                provided, a name is generated based on the class name and a
+                random suffix.
+            enforce (bool, optional): If False, only monitor the constraint
+                without adjusting the loss. Defaults to True.
+            rescale_factor (Number, optional): Factor to scale the
+                constraint-adjusted loss. Defaults to 1.5. Should be greater
+                than 1 to give weight to the constraint.
+
+        Raises:
+            TypeError: If a provided attribute has an incompatible type.
+            ValueError: If any tag in `tags` is not defined in the `descriptor`.
+
+        Note:
+            - If `rescale_factor <= 1`, a warning is issued.
+            - If `name` is not provided, a name is auto-generated, and a
+            warning is logged.
+        """
         # Init parent class
         super().__init__()
 
         # Type checking
-        validate_iterable("neurons", neurons, str)
+        validate_iterable("tags", tags, str)
         validate_type("name", name, str, allow_none=True)
-        validate_type("monitor_only", monitor_only, bool)
+        validate_type("enforce", enforce, bool)
         validate_type("rescale_factor", rescale_factor, Number)
 
         # Init object variables
-        self.neurons = neurons
+        self.tags = tags
         self.rescale_factor = rescale_factor
-        self.monitor_only = monitor_only
+        self.initial_rescale_factor = rescale_factor
+        self.enforce = enforce
 
         # Perform checks
         if rescale_factor <= 1:
@@ -135,120 +139,102 @@ class Constraint(ABC):
                 "will favor general loss over the constraint-adjusted loss. "
                 "Is this intended behavior? Normally, the rescale factor "
                 "should always be larger than 1.",
+                stacklevel=2,
             )
-        else:
-            self.rescale_factor = rescale_factor
 
         # If no constraint_name is set, generate one based
         # on the class name and a random suffix
         if name:
             self.name = name
         else:
-            random_suffix = "".join(
-                random.choices(string.ascii_uppercase + string.digits, k=6)
-            )
+            random_suffix = "".join(random.choices(string.ascii_uppercase + string.digits, k=6))
             self.name = f"{self.__class__.__name__}_{random_suffix}"
-            warnings.warn(
-                f"Name for constraint is not set. Using {self.name}.",
-            )
+            warnings.warn(f"Name for constraint is not set. Using {self.name}.", stacklevel=2)
 
-        # Infer layers from descriptor and neurons
+        # Infer layers from descriptor and tags
         self.layers = set()
-        for neuron in self.neurons:
-            if neuron not in self.descriptor.neuron_to_layer.keys():
+        for tag in self.tags:
+            if not self.descriptor.exists(tag):
                 raise ValueError(
-                    f"The neuron name {neuron} used with constraint "
+                    f"The tag {tag} used with constraint "
                     f"{self.name} is not defined in the descriptor. Please "
                     "add it to the correct layer using "
                     "descriptor.add('layer', ...)."
                 )
 
-            self.layers.add(self.descriptor.neuron_to_layer[neuron])
+            layer, _ = self.descriptor.location(tag)
+            self.layers.add(layer)
 
     @abstractmethod
-    def check_constraint(
-        self, prediction: dict[str, Tensor]
-    ) -> tuple[Tensor, int]:
-        """
-        Evaluates whether the given model predictions satisfy the constraint.
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Evaluates whether the given model predictions satisfy the constraint.
+
+        1 IS SATISFIED, 0 IS NOT SATISFIED
 
         Args:
-            prediction (dict[str, Tensor]): Model predictions for the neurons.
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
 
         Returns:
-            tuple[Tensor, int]: A tuple where the first element is a tensor
-            indicating whether the constraint is satisfied (with `True`
-            for satisfaction, `False` for non-satisfaction, and `torch.nan`
-            for irrelevant results), and the second element is an integer
-            value representing the number of relevant constraints.
+            tuple[Tensor, Tensor]: A tuple where the first element is a tensor of floats
+            indicating whether the constraint is satisfied (with value 1.0
+            for satisfaction, and 0.0 for non-satisfaction, and the second element is a tensor
+            mask that indicates the relevance of each sample (`True` for relevant
+            samples and `False` for irrelevant ones).
 
         Raises:
             NotImplementedError: If not implemented in a subclass.
         """
-
         raise NotImplementedError
 
     @abstractmethod
-    def calculate_direction(
-        self, prediction: dict[str, Tensor]
-    ) -> Dict[str, Tensor]:
-        """
-        Calculates adjustment directions for neurons to
-        better satisfy the constraint.
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Compute adjustment directions to better satisfy the constraint.
+
+        Given the model predictions, input batch, and context, this method calculates the direction
+        in which the predictions referenced by a tag should be adjusted to satisfy the constraint.
 
         Args:
-            prediction (dict[str, Tensor]): Model predictions for the neurons.
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
 
         Returns:
-            Dict[str, Tensor]: Dictionary mapping neuron layers to tensors
-            specifying the adjustment direction for each neuron.
+            dict[str, Tensor]: Dictionary mapping network layers to tensors that
+                specify the adjustment direction for each tag.
 
         Raises:
-            NotImplementedError: If not implemented in a subclass.
+            NotImplementedError: Must be implemented by subclasses.
         """
-
         raise NotImplementedError
 
 
 class ImplicationConstraint(Constraint):
-    """
-    Represents an implication constraint between two
-    constraints (head and body).
+    """Represents an implication constraint between two constraints (head and body).
 
     The implication constraint ensures that the `body` constraint only applies
     when the `head` constraint is satisfied. If the `head` constraint is not
     satisfied, the `body` constraint does not apply.
-
-    Args:
-        head (Constraint): The head of the implication. If this constraint
-            is satisfied, the body constraint must also be satisfied.
-        body (Constraint): The body of the implication. This constraint
-            is enforced only when the head constraint is satisfied.
-        name (str, optional): A unique name for the constraint. If not
-            provided, the name is generated in the format
-            "{body.name} if {head.name}". Defaults to None.
-        monitor_only (bool, optional): If True, the constraint is only
-            monitored without adjusting the loss. Defaults to False.
-        rescale_factor (Number, optional): The scaling factor for the
-            constraint-adjusted loss. Defaults to 1.5.
-
-    Raises:
-        TypeError: If a provided attribute has an incompatible type.
-
     """
 
     def __init__(
         self,
         head: Constraint,
         body: Constraint,
-        name=None,
-        monitor_only=False,
-        rescale_factor=1.5,
+        name: str = None,
     ):
-        """
-        Initializes an ImplicationConstraint instance.
-        """
+        """Initializes an ImplicationConstraint instance.
+
+        Uses `enforce` and `rescale_factor` from the body constraint.
+
+        Args:
+            head (Constraint): Constraint defining the head of the implication.
+            body (Constraint): Constraint defining the body of the implication.
+            name (str, optional): A unique name for the constraint. If not
+                provided, a name is generated based on the class name and a
+                random suffix.
+
+        Raises:
+            TypeError: If a provided attribute has an incompatible type.
 
+        """
         # Type checking
         validate_type("head", head, Constraint)
         validate_type("body", body, Constraint)
@@ -257,64 +243,82 @@ class ImplicationConstraint(Constraint):
         name = f"{body.name} if {head.name}"
 
         # Init parent class
-        super().__init__(
-            head.neurons | body.neurons,
-            name,
-            monitor_only,
-            rescale_factor,
-        )
+        super().__init__(head.tags | body.tags, name, body.enforce, body.rescale_factor)
 
         self.head = head
         self.body = body
 
-    def check_constraint(
-        self, prediction: dict[str, Tensor]
-    ) -> tuple[Tensor, int]:
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Check whether the implication constraint is satisfied.
+
+        Evaluates the `head` and `body` constraints. The `body` constraint
+        is enforced only if the `head` constraint is satisfied. If the
+        `head` constraint is not satisfied, the `body` constraint does not
+        affect the result.
+
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
 
+        Returns:
+            tuple[Tensor, Tensor]:
+                - result: Tensor indicating satisfaction of the implication
+                constraint (1 if satisfied, 0 otherwise).
+                - head_satisfaction: Tensor indicating satisfaction of the
+                head constraint alone.
+        """
         # Check satisfaction of head and body constraints
-        head_satisfaction, _ = self.head.check_constraint(prediction)
-        body_satisfaction, _ = self.body.check_constraint(prediction)
+        head_satisfaction, _ = self.head.check_constraint(data)
+        body_satisfaction, _ = self.body.check_constraint(data)
 
         # If head constraint is satisfied (returning 1),
         # the body constraint matters (and should return 0/1 based on body)
         # If head constraint is not satisfied (returning 0),
         # the body constraint does not apply (and should return 1)
-        result = logical_or(
-            logical_not(head_satisfaction), body_satisfaction
-        ).float()
+        result = logical_or(logical_not(head_satisfaction), body_satisfaction).float()
+
+        return result, head_satisfaction
 
-        return result, count_nonzero(head_satisfaction)
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Compute adjustment directions for tags to satisfy the constraint.
 
-    def calculate_direction(
-        self, prediction: dict[str, Tensor]
-    ) -> Dict[str, Tensor]:
+        Uses the `body` constraint directions as the update vector. Only
+        applies updates if the `head` constraint is satisfied. Currently,
+        this method only works for dense layers due to tag-to-index
+        translation limitations.
+
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
+
+        Returns:
+            dict[str, Tensor]: Dictionary mapping tags to tensors
+                specifying the adjustment direction for each tag.
+        """
         # NOTE currently only works for dense layers
-        # due to neuron to index translation
+        # due to tag to index translation
 
         # Use directions of constraint body as update vector
-        return self.body.calculate_direction(prediction)
+        return self.body.calculate_direction(data)
 
 
 class ScalarConstraint(Constraint):
-    """
-    A constraint that enforces scalar-based comparisons on a specific neuron.
+    """A constraint that enforces scalar-based comparisons on a specific tag.
 
-    This class ensures that the output of a specified neuron satisfies a scalar
+    This class ensures that the output of a specified tag satisfies a scalar
     comparison operation (e.g., less than, greater than, etc.). It uses a
     comparator function to validate the condition and calculates adjustment
     directions accordingly.
 
     Args:
-        operand (Union[str, Transformation]): Name of the neuron or a
+        operand (Union[str, Transformation]): Name of the tag or a
             transformation to apply.
         comparator (Callable[[Tensor, Number], Tensor]): A comparison
             function (e.g., `torch.ge`, `torch.lt`).
         scalar (Number): The scalar value to compare against.
         name (str, optional): A unique name for the constraint. If not
             provided, a name is auto-generated in the format
-            "<neuron_name> <comparator> <scalar>".
-        monitor_only (bool, optional): If True, only monitor the constraint
-            without adjusting the loss. Defaults to False.
+            "<tag> <comparator> <scalar>".
+        enforce (bool, optional): If False, only monitor the constraint
+            without adjusting the loss. Defaults to True.
         rescale_factor (Number, optional): Factor to scale the
             constraint-adjusted loss. Defaults to 1.5.
 
@@ -322,87 +326,120 @@ class ScalarConstraint(Constraint):
         TypeError: If a provided attribute has an incompatible type.
 
     Notes:
-        - The `neuron_name` must be defined in the `descriptor` mapping.
-        - The constraint name is composed using the neuron name,
-          comparator, and scalar value.
+        - The `tag` must be defined in the `descriptor` mapping.
+        - The constraint name is composed using the tag, comparator, and scalar value.
 
     """
 
     def __init__(
         self,
-        operand: Union[str, Transformation],
+        operand: str | Transformation,
         comparator: Callable[[Tensor, Number], Tensor],
         scalar: Number,
         name: str = None,
-        monitor_only: bool = False,
+        enforce: bool = True,
         rescale_factor: Number = 1.5,
     ) -> None:
-        """
-        Initializes a ScalarConstraint instance.
-        """
+        """Initializes a ScalarConstraint instance.
+
+        Args:
+            operand (Union[str, Transformation]): Function that needs to be
+                performed on the network variables before applying the
+                constraint.
+            comparator (Callable[[Tensor, Number], Tensor]): Comparison
+                operator used in the constraint. Supported types are
+                {torch.lt, torch.le, torch.st, torch.se}.
+            scalar (Number): Constant to compare the variable to.
+            name (str, optional): A unique name for the constraint. If not
+                provided, a name is generated based on the class name and a
+                random suffix.
+            enforce (bool, optional): If False, only monitor the constraint
+                without adjusting the loss. Defaults to True.
+            rescale_factor (Number, optional): Factor to scale the
+                constraint-adjusted loss. Defaults to 1.5. Should be greater
+                than 1 to give weight to the constraint.
+
+        Raises:
+            TypeError: If a provided attribute has an incompatible type.
 
+        Notes:
+            - The `tag` must be defined in the `descriptor` mapping.
+            - The constraint name is composed using the tag, comparator, and scalar value.
+        """
         # Type checking
         validate_type("operand", operand, (str, Transformation))
         validate_comparator_pytorch("comparator", comparator)
-        validate_comparator_pytorch("comparator", comparator)
         validate_type("scalar", scalar, Number)
 
-        # If transformation is provided, get neuron name,
-        # else use IdentityTransformation
+        # If transformation is provided, get tag name, else use IdentityTransformation
         if isinstance(operand, Transformation):
-            neuron_name = operand.neuron_name
+            tag = operand.tag
             transformation = operand
         else:
-            neuron_name = operand
-            transformation = IdentityTransformation(neuron_name)
+            tag = operand
+            transformation = IdentityTransformation(tag)
 
         # Compose constraint name
-        name = f"{neuron_name} {comparator.__name__} {str(scalar)}"
+        name = f"{tag} {comparator.__name__} {str(scalar)}"
 
         # Init parent class
-        super().__init__({neuron_name}, name, monitor_only, rescale_factor)
+        super().__init__({tag}, name, enforce, rescale_factor)
 
         # Init variables
+        self.tag = tag
         self.comparator = comparator
         self.scalar = scalar
         self.transformation = transformation
 
-        # Get layer name and feature index from neuron_name
-        self.layer = self.descriptor.neuron_to_layer[neuron_name]
-        self.index = self.descriptor.neuron_to_index[neuron_name]
-
         # Calculate directions based on constraint operator
         if self.comparator in [lt, le]:
-            self.direction = -1
-        elif self.comparator in [gt, ge]:
             self.direction = 1
+        elif self.comparator in [gt, ge]:
+            self.direction = -1
+
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Check if the scalar constraint is satisfied for a given tag.
 
-    def check_constraint(
-        self, prediction: dict[str, Tensor]
-    ) -> tuple[Tensor, int]:
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
 
+        Returns:
+            tuple[Tensor, Tensor]:
+                - result: Tensor indicating whether the tag satisfies the constraint.
+                - ones_like(result): Tensor of ones with same shape as `result`.
+        """
         # Select relevant columns
-        selection = prediction[self.layer][:, self.index]
+        selection = self.descriptor.select(self.tag, data)
 
         # Apply transformation
         selection = self.transformation(selection)
 
         # Calculate current constraint result
         result = self.comparator(selection, self.scalar).float()
-        return result, numel(result)
+        return result, ones_like(result)
+
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Compute adjustment directions to satisfy the scalar constraint.
+
+        Only works for dense layers due to tag-to-index translation.
+
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
 
-    def calculate_direction(
-        self, prediction: dict[str, Tensor]
-    ) -> Dict[str, Tensor]:
+        Returns:
+            dict[str, Tensor]: Dictionary mapping layers to tensors specifying
+                the adjustment direction for each tag.
+        """
         # NOTE currently only works for dense layers due
-        # to neuron to index translation
+        # to tag to index translation
 
         output = {}
 
         for layer in self.layers:
-            output[layer] = zeros_like(prediction[layer][0], device=self.device)
+            output[layer] = zeros_like(data[layer][0], device=self.device)
 
-        output[self.layer][self.index] = self.direction
+        layer, index = self.descriptor.location(self.tag)
+        output[layer][index] = self.direction
 
         for layer in self.layers:
             output[layer] = normalize(reshape(output[layer], [1, -1]), dim=1)
@@ -411,26 +448,24 @@ class ScalarConstraint(Constraint):
 
 
 class BinaryConstraint(Constraint):
-    """
-    A constraint that enforces a binary comparison between two neurons.
+    """A constraint that enforces a binary comparison between two tags.
 
-    This class ensures that the output of one neuron satisfies a comparison
-    operation with the output of another neuron
-    (e.g., less than, greater than, etc.). It uses a comparator function to
-    validate the condition and calculates adjustment directions accordingly.
+    This class ensures that the output of one tag satisfies a comparison
+    operation with the output of another tag (e.g., less than, greater than, etc.).
+    It uses a comparator function to validate the condition and calculates adjustment directions accordingly.
 
     Args:
         operand_left (Union[str, Transformation]): Name of the left
-            neuron or a transformation to apply.
+            tag or a transformation to apply.
         comparator (Callable[[Tensor, Number], Tensor]): A comparison
             function (e.g., `torch.ge`, `torch.lt`).
         operand_right (Union[str, Transformation]): Name of the right
-            neuron or a transformation to apply.
+            tag or a transformation to apply.
         name (str, optional): A unique name for the constraint. If not
             provided, a name is auto-generated in the format
-            "<neuron_name_left> <comparator> <neuron_name_right>".
-        monitor_only (bool, optional): If True, only monitor the constraint
-            without adjusting the loss. Defaults to False.
+            "<operand_left> <comparator> <operand_right>".
+        enforce (bool, optional): If False, only monitor the constraint
+            without adjusting the loss. Defaults to True.
         rescale_factor (Number, optional): Factor to scale the
             constraint-adjusted loss. Defaults to 1.5.
 
@@ -438,84 +473,107 @@ class BinaryConstraint(Constraint):
         TypeError: If a provided attribute has an incompatible type.
 
     Notes:
-        - The neuron names must be defined in the `descriptor` mapping.
-        - The constraint name is composed using the left neuron name,
-          comparator, and right neuron name.
+        - The tags must be defined in the `descriptor` mapping.
+        - The constraint name is composed using the left tag, comparator, and right tag.
 
     """
 
     def __init__(
         self,
-        operand_left: Union[str, Transformation],
+        operand_left: str | Transformation,
         comparator: Callable[[Tensor, Number], Tensor],
-        operand_right: Union[str, Transformation],
+        operand_right: str | Transformation,
         name: str = None,
-        monitor_only: bool = False,
+        enforce: bool = True,
         rescale_factor: Number = 1.5,
     ) -> None:
-        """
-        Initializes a BinaryConstraint instance.
-        """
+        """Initializes a BinaryConstraint instance.
 
+        Args:
+            operand_left (Union[str, Transformation]): Name of the left
+                tag or a transformation to apply.
+            comparator (Callable[[Tensor, Number], Tensor]): A comparison
+                function (e.g., `torch.ge`, `torch.lt`).
+            operand_right (Union[str, Transformation]): Name of the right
+                tag or a transformation to apply.
+            name (str, optional): A unique name for the constraint. If not
+                provided, a name is auto-generated in the format
+                "<operand_left> <comparator> <operand_right>".
+            enforce (bool, optional): If False, only monitor the constraint
+                without adjusting the loss. Defaults to True.
+            rescale_factor (Number, optional): Factor to scale the
+                constraint-adjusted loss. Defaults to 1.5.
+
+        Raises:
+            TypeError: If a provided attribute has an incompatible type.
+
+        Notes:
+            - The tags must be defined in the `descriptor` mapping.
+            - The constraint name is composed using the left tag,
+              comparator, and right tag.
+        """
         # Type checking
         validate_type("operand_left", operand_left, (str, Transformation))
         validate_comparator_pytorch("comparator", comparator)
         validate_comparator_pytorch("comparator", comparator)
         validate_type("operand_right", operand_right, (str, Transformation))
 
-        # If transformation is provided, get neuron name,
-        # else use IdentityTransformation
+        # If transformation is provided, get tag name, else use IdentityTransformation
         if isinstance(operand_left, Transformation):
-            neuron_name_left = operand_left.neuron_name
+            tag_left = operand_left.tag
             transformation_left = operand_left
         else:
-            neuron_name_left = operand_left
-            transformation_left = IdentityTransformation(neuron_name_left)
+            tag_left = operand_left
+            transformation_left = IdentityTransformation(tag_left)
 
         if isinstance(operand_right, Transformation):
-            neuron_name_right = operand_right.neuron_name
+            tag_right = operand_right.tag
             transformation_right = operand_right
         else:
-            neuron_name_right = operand_right
-            transformation_right = IdentityTransformation(neuron_name_right)
+            tag_right = operand_right
+            transformation_right = IdentityTransformation(tag_right)
 
         # Compose constraint name
-        name = f"{neuron_name_left} {comparator.__name__} {neuron_name_right}"
+        name = f"{tag_left} {comparator.__name__} {tag_right}"
 
         # Init parent class
-        super().__init__(
-            {neuron_name_left, neuron_name_right},
-            name,
-            monitor_only,
-            rescale_factor,
-        )
+        super().__init__({tag_left, tag_right}, name, enforce, rescale_factor)
 
         # Init variables
         self.comparator = comparator
+        self.tag_left = tag_left
+        self.tag_right = tag_right
         self.transformation_left = transformation_left
         self.transformation_right = transformation_right
 
-        # Get layer name and feature index from neuron_name
-        self.layer_left = self.descriptor.neuron_to_layer[neuron_name_left]
-        self.layer_right = self.descriptor.neuron_to_layer[neuron_name_right]
-        self.index_left = self.descriptor.neuron_to_index[neuron_name_left]
-        self.index_right = self.descriptor.neuron_to_index[neuron_name_right]
-
         # Calculate directions based on constraint operator
         if self.comparator in [lt, le]:
-            self.direction_left = -1
-            self.direction_right = 1
-        else:
             self.direction_left = 1
             self.direction_right = -1
+        else:
+            self.direction_left = -1
+            self.direction_right = 1
+
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Evaluate whether the binary constraint is satisfied for the current predictions.
 
-    def check_constraint(
-        self, prediction: dict[str, Tensor]
-    ) -> tuple[Tensor, int]:
+        The constraint compares the outputs of two tags using the specified
+        comparator function. A result of `1` indicates the constraint is satisfied
+        for a sample, and `0` indicates it is violated.
+
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
 
+        Returns:
+            tuple[Tensor, Tensor]:
+                - result (Tensor): Binary tensor indicating constraint satisfaction
+                (1 for satisfied, 0 for violated) for each sample.
+                - mask (Tensor): Tensor of ones with the same shape as `result`,
+                used for constraint aggregation.
+        """
         # Select relevant columns
-        selection_left = prediction[self.layer_left][:, self.index_left]
-        selection_right = prediction[self.layer_right][:, self.index_right]
+        selection_left = self.descriptor.select(self.tag_left, data)
+        selection_right = self.descriptor.select(self.tag_right, data)
 
         # Apply transformations
         selection_left = self.transformation_left(selection_left)
@@ -523,21 +581,34 @@ class BinaryConstraint(Constraint):
 
         result = self.comparator(selection_left, selection_right).float()
 
-        return result, numel(result)
+        return result, ones_like(result)
+
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Compute adjustment directions for the tags involved in the binary constraint.
+
+        The returned directions indicate how to adjust each tag's output to
+        satisfy the constraint. Only currently supported for dense layers.
 
-    def calculate_direction(
-        self, prediction: dict[str, Tensor]
-    ) -> Dict[str, Tensor]:
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
+
+        Returns:
+            dict[str, Tensor]: A mapping from layer names to tensors specifying
+            the normalized adjustment directions for each tag involved in the
+            constraint.
+        """
         # NOTE currently only works for dense layers due
-        # to neuron to index translation
+        # to tag to index translation
 
         output = {}
 
         for layer in self.layers:
-            output[layer] = zeros_like(prediction[layer][0], device=self.device)
+            output[layer] = zeros_like(data[layer][0], device=self.device)
 
-        output[self.layer_left][self.index_left] = self.direction_left
-        output[self.layer_right][self.index_right] = self.direction_right
+        layer_left, index_left = self.descriptor.location(self.tag_left)
+        layer_right, index_right = self.descriptor.location(self.tag_right)
+        output[layer_left][index_left] = self.direction_left
+        output[layer_right][index_right] = self.direction_right
 
         for layer in self.layers:
             output[layer] = normalize(reshape(output[layer], [1, -1]), dim=1)
@@ -546,142 +617,119 @@ class BinaryConstraint(Constraint):
 
 
 class SumConstraint(Constraint):
-    """
-    A constraint that enforces a weighted summation comparison
-    between two groups of neurons.
+    """A constraint that enforces a weighted summation comparison between two groups of tags.
 
     This class evaluates whether the weighted sum of outputs from one set of
-    neurons satisfies a comparison operation with the weighted sum of
-    outputs from another set of neurons.
-
-    Args:
-        operands_left (list[Union[str, Transformation]]): List of neuron
-            names or transformations on the left side.
-        comparator (Callable[[Tensor, Number], Tensor]): A comparison
-            function for the constraint.
-        operands_right (list[Union[str, Transformation]]): List of neuron
-            names or transformations on the right side.
-        weights_left (list[Number], optional): Weights for the left neurons.
-            Defaults to None.
-        weights_right (list[Number], optional): Weights for the right
-            neurons. Defaults to None.
-        name (str, optional): Unique name for the constraint.
-            If None, it's auto-generated. Defaults to None.
-        monitor_only (bool, optional): If True, only monitor the constraint
-            without adjusting the loss. Defaults to False.
-        rescale_factor (Number, optional): Factor to scale the
-            constraint-adjusted loss. Defaults to 1.5.
-
-    Raises:
-        TypeError: If a provided attribute has an incompatible type.
-        ValueError: If the dimensions of neuron names and weights mismatch.
-
+    tags satisfies a comparison operation with the weighted sum of
+    outputs from another set of tags.
     """
 
     def __init__(
         self,
-        operands_left: list[Union[str, Transformation]],
+        operands_left: list[str | Transformation],
         comparator: Callable[[Tensor, Number], Tensor],
-        operands_right: list[Union[str, Transformation]],
+        operands_right: list[str | Transformation],
         weights_left: list[Number] = None,
         weights_right: list[Number] = None,
         name: str = None,
-        monitor_only: bool = False,
+        enforce: bool = True,
         rescale_factor: Number = 1.5,
     ) -> None:
-        """
-        Initializes the SumConstraint.
-        """
+        """Initializes the SumConstraint.
 
+        Args:
+            operands_left (list[Union[str, Transformation]]): List of tags
+                or transformations on the left side.
+            comparator (Callable[[Tensor, Number], Tensor]): A comparison
+                function for the constraint.
+            operands_right (list[Union[str, Transformation]]): List of tags
+                or transformations on the right side.
+            weights_left (list[Number], optional): Weights for the left
+                tags. Defaults to None.
+            weights_right (list[Number], optional): Weights for the right
+                tags. Defaults to None.
+            name (str, optional): Unique name for the constraint.
+                If None, it's auto-generated. Defaults to None.
+            enforce (bool, optional): If False, only monitor the constraint
+                without adjusting the loss. Defaults to True.
+            rescale_factor (Number, optional): Factor to scale the
+                constraint-adjusted loss. Defaults to 1.5.
+
+        Raises:
+            TypeError: If a provided attribute has an incompatible type.
+            ValueError: If the dimensions of tags and weights mismatch.
+        """
         # Type checking
         validate_iterable("operands_left", operands_left, (str, Transformation))
         validate_comparator_pytorch("comparator", comparator)
         validate_comparator_pytorch("comparator", comparator)
-        validate_iterable(
-            "operands_right", operands_right, (str, Transformation)
-        )
+        validate_iterable("operands_right", operands_right, (str, Transformation))
         validate_iterable("weights_left", weights_left, Number, allow_none=True)
-        validate_iterable(
-            "weights_right", weights_right, Number, allow_none=True
-        )
+        validate_iterable("weights_right", weights_right, Number, allow_none=True)
 
-        # If transformation is provided, get neuron name,
-        # else use IdentityTransformation
-        neuron_names_left: list[str] = []
+        # If transformation is provided, get tag, else use IdentityTransformation
+        tags_left: list[str] = []
         transformations_left: list[Transformation] = []
         for operand_left in operands_left:
             if isinstance(operand_left, Transformation):
-                neuron_name_left = operand_left.neuron_name
-                neuron_names_left.append(neuron_name_left)
+                tag_left = operand_left.tag
+                tags_left.append(tag_left)
                 transformations_left.append(operand_left)
             else:
-                neuron_name_left = operand_left
-                neuron_names_left.append(neuron_name_left)
-                transformations_left.append(
-                    IdentityTransformation(neuron_name_left)
-                )
+                tag_left = operand_left
+                tags_left.append(tag_left)
+                transformations_left.append(IdentityTransformation(tag_left))
 
-        neuron_names_right: list[str] = []
+        tags_right: list[str] = []
         transformations_right: list[Transformation] = []
         for operand_right in operands_right:
             if isinstance(operand_right, Transformation):
-                neuron_name_right = operand_right.neuron_name
-                neuron_names_right.append(neuron_name_right)
+                tag_right = operand_right.tag
+                tags_right.append(tag_right)
                 transformations_right.append(operand_right)
             else:
-                neuron_name_right = operand_right
-                neuron_names_right.append(neuron_name_right)
-                transformations_right.append(
-                    IdentityTransformation(neuron_name_right)
-                )
+                tag_right = operand_right
+                tags_right.append(tag_right)
+                transformations_right.append(IdentityTransformation(tag_right))
 
         # Compose constraint name
-        w_left = weights_left or [""] * len(neuron_names_left)
-        w_right = weights_right or [""] * len(neuron_names_right)
-        left_expr = " + ".join(
-            f"{w}{n}" for w, n in zip(w_left, neuron_names_left)
-        )
-        right_expr = " + ".join(
-            f"{w}{n}" for w, n in zip(w_right, neuron_names_right)
-        )
+        w_left = weights_left or [""] * len(tags_left)
+        w_right = weights_right or [""] * len(tags_right)
+        left_expr = " + ".join(f"{w}{n}" for w, n in zip(w_left, tags_left, strict=False))
+        right_expr = " + ".join(f"{w}{n}" for w, n in zip(w_right, tags_right, strict=False))
         comparator_name = comparator.__name__
         name = f"{left_expr} {comparator_name} {right_expr}"
 
         # Init parent class
-        neuron_names = set(neuron_names_left) | set(neuron_names_right)
-        super().__init__(neuron_names, name, monitor_only, rescale_factor)
+        tags = set(tags_left) | set(tags_right)
+        super().__init__(tags, name, enforce, rescale_factor)
 
         # Init variables
         self.comparator = comparator
-        self.neuron_names_left = neuron_names_left
-        self.neuron_names_right = neuron_names_right
+        self.tags_left = tags_left
+        self.tags_right = tags_right
         self.transformations_left = transformations_left
         self.transformations_right = transformations_right
 
-        # If feature list dimensions don't match
-        # weight list dimensions, raise error
-        if weights_left and (len(neuron_names_left) != len(weights_left)):
+        # If feature list dimensions don't match weight list dimensions, raise error
+        if weights_left and (len(tags_left) != len(weights_left)):
             raise ValueError(
-                "The dimensions of neuron_names_left don't match with the "
-                "dimensions of weights_left."
+                "The dimensions of tags_left don't match with the dimensions of weights_left."
             )
-        if weights_right and (len(neuron_names_right) != len(weights_right)):
+        if weights_right and (len(tags_right) != len(weights_right)):
             raise ValueError(
-                "The dimensions of neuron_names_right don't match with the "
-                "dimensions of weights_right."
+                "The dimensions of tags_right don't match with the dimensions of weights_right."
             )
 
         # If weights are provided for summation, transform them to Tensors
         if weights_left:
             self.weights_left = tensor(weights_left, device=self.device)
         else:
-            self.weights_left = ones(len(neuron_names_left), device=self.device)
+            self.weights_left = ones(len(tags_left), device=self.device)
         if weights_right:
             self.weights_right = tensor(weights_right, device=self.device)
         else:
-            self.weights_right = ones(
-                len(neuron_names_right), device=self.device
-            )
+            self.weights_right = ones(len(tags_right), device=self.device)
 
         # Calculate directions based on constraint operator
         if self.comparator in [lt, le]:
@@ -691,80 +739,82 @@ class SumConstraint(Constraint):
             self.direction_left = 1
             self.direction_right = -1
 
-    def check_constraint(
-        self, prediction: dict[str, Tensor]
-    ) -> tuple[Tensor, int]:
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Evaluate whether the weighted sum constraint is satisfied.
+
+        Computes the weighted sum of outputs from the left and right tags,
+        applies the specified comparator function, and returns a binary result for
+        each sample.
+
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
+
+        Returns:
+            tuple[Tensor, Tensor]:
+                - result (Tensor): Binary tensor indicating whether the constraint
+                is satisfied (1) or violated (0) for each sample.
+                - mask (Tensor): Tensor of ones, used for constraint aggregation.
+        """
 
         def compute_weighted_sum(
-            neuron_names: list[str],
+            tags: list[str],
             transformations: list[Transformation],
-            weights: tensor,
-        ) -> tensor:
-            layers = [
-                self.descriptor.neuron_to_layer[neuron_name]
-                for neuron_name in neuron_names
-            ]
-            indices = [
-                self.descriptor.neuron_to_index[neuron_name]
-                for neuron_name in neuron_names
-            ]
-
-            # Select relevant column
-            selections = [
-                prediction[layer][:, index]
-                for layer, index in zip(layers, indices)
-            ]
+            weights: Tensor,
+        ) -> Tensor:
+            # Select relevant columns
+            selections = [self.descriptor.select(tag, data) for tag in tags]
 
             # Apply transformations
             results = []
-            for transformation, selection in zip(transformations, selections):
+            for transformation, selection in zip(transformations, selections, strict=False):
                 results.append(transformation(selection))
 
-            # Extract predictions for all neurons and apply weights in bulk
-            predictions = stack(
-                results,
-                dim=1,
-            )
+            # Extract predictions for all tags and apply weights in bulk
+            predictions = stack(results)
 
             # Calculate weighted sum
-            return (predictions * weights.unsqueeze(0)).sum(dim=1)
+            return (predictions * weights.view(-1, 1, 1)).sum(dim=0)
 
         # Compute weighted sums
         weighted_sum_left = compute_weighted_sum(
-            self.neuron_names_left,
-            self.transformations_left,
-            self.weights_left,
+            self.tags_left, self.transformations_left, self.weights_left
         )
         weighted_sum_right = compute_weighted_sum(
-            self.neuron_names_right,
-            self.transformations_right,
-            self.weights_right,
+            self.tags_right, self.transformations_right, self.weights_right
         )
 
         # Apply the comparator and calculate the result
         result = self.comparator(weighted_sum_left, weighted_sum_right).float()
 
-        return result, numel(result)
+        return result, ones_like(result)
+
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Compute adjustment directions for tags involved in the weighted sum constraint.
+
+        The directions indicate how to adjust each tag's output to satisfy the
+        constraint. Only dense layers are currently supported.
+
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
 
-    def calculate_direction(
-        self, prediction: dict[str, Tensor]
-    ) -> Dict[str, Tensor]:
+        Returns:
+            dict[str, Tensor]: Mapping from layer names to normalized tensors
+            specifying adjustment directions for each tag involved in the constraint.
+        """
         # NOTE currently only works for dense layers
-        # due to neuron to index translation
+        # due to tag to index translation
 
         output = {}
 
         for layer in self.layers:
-            output[layer] = zeros_like(prediction[layer][0], device=self.device)
+            output[layer] = zeros_like(data[layer][0], device=self.device)
 
-        for neuron_name_left in self.neuron_names_left:
-            layer = self.descriptor.neuron_to_layer[neuron_name_left]
-            index = self.descriptor.neuron_to_index[neuron_name_left]
+        for tag_left in self.tags_left:
+            layer, index = self.descriptor.location(tag_left)
             output[layer][index] = self.direction_left
 
-        for neuron_name_right in self.neuron_names_right:
-            layer = self.descriptor.neuron_to_layer[neuron_name_right]
-            index = self.descriptor.neuron_to_index[neuron_name_right]
+        for tag_right in self.tags_right:
+            layer, index = self.descriptor.location(tag_right)
             output[layer][index] = self.direction_right
 
         for layer in self.layers:
@@ -773,134 +823,434 @@ class SumConstraint(Constraint):
         return output
 
 
-class PythagoreanIdentityConstraint(Constraint):
+class MonotonicityConstraint(Constraint):
+    """Constraint that enforces a monotonic relationship between two tags.
+
+    This constraint ensures that the activations of a prediction tag (`tag_prediction`)
+    are monotonically ascending or descending with respect to a target tag (`tag_reference`).
     """
-    A constraint that enforces the Pythagorean identity: a² + b² ≈ 1,
-    where `a` and `b` are neurons or transformations.
 
-    This constraint checks that the sum of the squares of two specified
-    neurons (or their transformations) is approximately equal to 1.
-    The constraint is evaluated using relative and absolute
-    tolerance (`rtol` and `atol`) and is applied during the forward pass.
+    def __init__(
+        self,
+        tag_prediction: str,
+        tag_reference: str,
+        rescale_factor_lower: float = 1.5,
+        rescale_factor_upper: float = 1.75,
+        stable: bool = True,
+        direction: Literal["ascending", "descending"] = "ascending",
+        name: str = None,
+        enforce: bool = True,
+    ):
+        """Constraint that enforces monotonicity on a predicted output.
 
-    Args:
-        a (Union[str, Transformation]): The first input, either a
-            neuron name (str) or a Transformation.
-        b (Union[str, Transformation]): The second input, either a
-            neuron name (str) or a Transformation.
-        rtol (float, optional): The relative tolerance for the
-            comparison (default is 0.00001).
-        atol (float, optional): The absolute tolerance for the
-            comparison (default is 1e-8).
-        name (str, optional): The name of the constraint
-            (default is None, and it is generated automatically).
-        monitor_only (bool, optional): Flag indicating whether the
-            constraint is only for monitoring (default is False).
-        rescale_factor (Number, optional): A factor used for
-            rescaling (default is 1.5).
+        This constraint ensures that the activations of a prediction tag (`tag_prediction`)
+        are monotonically ascending or descending with respect to a target tag (`tag_reference`).
 
-    Raises:
-        TypeError: If a provided attribute has an incompatible type.
+        Args:
+            tag_prediction (str): Name of the tag whose activations should follow the monotonic relationship.
+            tag_reference (str): Name of the tag that acts as the monotonic reference.
+            rescale_factor_lower (float, optional): Lower bound for rescaling rank differences. Defaults to 1.5.
+            rescale_factor_upper (float, optional): Upper bound for rescaling rank differences. Defaults to 1.75.
+            stable (bool, optional): Whether to use stable sorting when ranking. Defaults to True.
+            direction (str, optional): Direction of monotonicity to enforce, either 'ascending' or 'descending'. Defaults to 'ascending'.
+            name (str, optional): Custom name for the constraint. If None, a descriptive name is auto-generated.
+            enforce (bool, optional): If False, the constraint is only monitored (not enforced). Defaults to True.
+        """
+        # Type checking
+        validate_type("rescale_factor_lower", rescale_factor_lower, float)
+        validate_type("rescale_factor_upper", rescale_factor_upper, float)
+        validate_type("stable", stable, bool)
+        validate_type("direction", direction, str)
+
+        # Compose constraint name
+        if name is None:
+            name = f"{tag_prediction} monotonically {direction} by {tag_reference}"
+
+        # Init parent class
+        super().__init__({tag_prediction}, name, enforce, 1.0)
+
+        # Init variables
+        self.tag_prediction = tag_prediction
+        self.tag_reference = tag_reference
+        self.rescale_factor_lower = rescale_factor_lower
+        self.rescale_factor_upper = rescale_factor_upper
+        self.stable = stable
+        self.descending = direction == "descending"
+
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Evaluate whether the monotonicity constraint is satisfied."""
+        # Select relevant columns
+        preds = self.descriptor.select(self.tag_prediction, data)
+        targets = self.descriptor.select(self.tag_reference, data)
+
+        # Utility: convert values -> ranks (0 ... num_features-1)
+        def compute_ranks(x: Tensor, descending: bool) -> Tensor:
+            return argsort(
+                argsort(x, descending=descending, stable=self.stable, dim=0),
+                descending=False,
+                stable=self.stable,
+                dim=0,
+            )
 
+        # Compute predicted and target ranks
+        pred_ranks = compute_ranks(preds, descending=self.descending)
+        target_ranks = compute_ranks(targets, descending=False)
+
+        # Rank difference
+        rank_diff = pred_ranks - target_ranks
+
+        # Rescale differences into [rescale_factor_lower, rescale_factor_upper]
+        batch_size = preds.shape[0]
+        invert_direction = -1 if self.descending else 1
+        self.compared_rankings = (
+            (rank_diff / batch_size) * (self.rescale_factor_upper - self.rescale_factor_lower)
+            + self.rescale_factor_lower * sign(rank_diff)
+        ) * invert_direction
+
+        # Calculate satisfaction
+        incorrect_rankings = eq(self.compared_rankings, 0).float()
+
+        return incorrect_rankings, ones_like(incorrect_rankings)
+
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Calculates ranking adjustments for monotonicity enforcement."""
+        layer, _ = self.descriptor.location(self.tag_prediction)
+        return {layer: self.compared_rankings}
+
+
+class GroupedMonotonicityConstraint(MonotonicityConstraint):
+    """Constraint that enforces a monotonic relationship between two tags.
+
+    This constraint ensures that the activations of a prediction tag (`tag_prediction`)
+    are monotonically ascending or descending with respect to a target tag (`tag_reference`).
     """
 
     def __init__(
         self,
-        a: Union[str, Transformation],
-        b: Union[str, Transformation],
-        rtol: float = 0.00001,
-        atol: float = 1e-8,
+        tag_prediction: str,
+        tag_reference: str,
+        tag_group_identifier: str,
+        rescale_factor_lower: float = 1.5,
+        rescale_factor_upper: float = 1.75,
+        stable: bool = True,
+        direction: Literal["ascending", "descending"] = "ascending",
+        name: str = None,
+        enforce: bool = True,
+    ):
+        """Constraint that enforces monotonicity on a predicted output.
+
+        This constraint ensures that the activations of a prediction tag (`tag_prediction`)
+        are monotonically ascending or descending with respect to a target tag (`tag_reference`).
+
+        Args:
+            tag_prediction (str): Name of the tag whose activations should follow the monotonic relationship.
+            tag_reference (str): Name of the tag that acts as the monotonic reference.
+            tag_group_identifier (str): Name of the tag that identifies groups for separate monotonicity enforcement.
+            rescale_factor_lower (float, optional): Lower bound for rescaling rank differences. Defaults to 1.5.
+            rescale_factor_upper (float, optional): Upper bound for rescaling rank differences. Defaults to 1.75.
+            stable (bool, optional): Whether to use stable sorting when ranking. Defaults to True.
+            direction (str, optional): Direction of monotonicity to enforce, either 'ascending' or 'descending'. Defaults to 'ascending'.
+            name (str, optional): Custom name for the constraint. If None, a descriptive name is auto-generated.
+            enforce (bool, optional): If False, the constraint is only monitored (not enforced). Defaults to True.
+        """
+        # Compose constraint name
+        if name is None:
+            name = f"{tag_prediction} for each {tag_group_identifier} monotonically {direction} by {tag_reference}"
+
+        # Init parent class
+        super().__init__(
+            tag_prediction=tag_prediction,
+            tag_reference=tag_reference,
+            rescale_factor_lower=rescale_factor_lower,
+            rescale_factor_upper=rescale_factor_upper,
+            stable=stable,
+            direction=direction,
+            name=name,
+            enforce=enforce,
+        )
+
+        # Init variables
+        self.tag_prediction = tag_prediction
+        self.tag_reference = tag_reference
+        self.tag_group_identifier = tag_group_identifier
+
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Evaluate whether the monotonicity constraint is satisfied."""
+        # Select group identifiers and convert to unique list
+        group_identifiers = self.descriptor.select(self.tag_group_identifier, data)
+        unique_group_identifiers = unique(group_identifiers, sorted=False).tolist()
+
+        # Initialize checks and directions
+        checks = zeros_like(group_identifiers, device=self.device)
+        self.directions = zeros_like(group_identifiers, device=self.device)
+
+        # Get prediction and target keys
+        preds_key, _ = self.descriptor.location(self.tag_prediction)
+        targets_key, _ = self.descriptor.location(self.tag_reference)
+
+        for group_identifier in unique_group_identifiers:
+            # Create mask for the samples in this group
+            group_mask = (group_identifiers == group_identifier).squeeze(1)
+
+            # Create mini-batch for the group
+            group_data = {
+                preds_key: data[preds_key][group_mask],
+                targets_key: data[targets_key][group_mask],
+            }
+
+            # Call super on the mini-batch
+            checks[group_mask], _ = super().check_constraint(group_data)
+            self.directions[group_mask] = self.compared_rankings
+
+        return checks, ones_like(checks)
+
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Calculates ranking adjustments for monotonicity enforcement."""
+        layer, _ = self.descriptor.location(self.tag_prediction)
+        return {layer: self.directions}
+
+
+class ANDConstraint(Constraint):
+    """A composite constraint that enforces the logical AND of multiple constraints.
+
+    This class combines multiple sub-constraints and evaluates them jointly:
+
+    * The satisfaction of the AND constraint is `True` only if all sub-constraints
+    are satisfied (elementwise logical AND).
+    * The corrective direction is computed by weighting each sub-constraint's
+    direction with its satisfaction mask and summing across all sub-constraints.
+    """
+
+    def __init__(
+        self,
+        *constraints: Constraint,
         name: str = None,
         monitor_only: bool = False,
         rescale_factor: Number = 1.5,
     ) -> None:
+        """A composite constraint that enforces the logical AND of multiple constraints.
+
+        This class combines multiple sub-constraints and evaluates them jointly:
+
+        * The satisfaction of the AND constraint is `True` only if all sub-constraints
+        are satisfied (elementwise logical AND).
+        * The corrective direction is computed by weighting each sub-constraint's
+        direction with its satisfaction mask and summing across all sub-constraints.
+
+        Args:
+            *constraints (Constraint): One or more `Constraint` instances to be combined.
+            name (str, optional): A custom name for this constraint. If not provided,
+                the name will be composed from the sub-constraint names joined with
+                " AND ".
+            monitor_only (bool, optional): If True, the constraint will be monitored
+                but not enforced. Defaults to False.
+            rescale_factor (Number, optional): A scaling factor applied when rescaling
+                corrections. Defaults to 1.5.
+
+        Attributes:
+            constraints (tuple[Constraint, ...]): The sub-constraints being combined.
+            neurons (set): The union of neurons referenced by the sub-constraints.
+            name (str): The name of the constraint (composed or custom).
         """
-        Initialize the PythagoreanIdentityConstraint.
+        # Type checking
+        validate_iterable("constraints", constraints, Constraint)
+
+        # Compose constraint name
+        if not name:
+            name = " AND ".join([constraint.name for constraint in constraints])
+
+        # Init parent class
+        super().__init__(
+            set().union(*(constraint.tags for constraint in constraints)),
+            name,
+            monitor_only,
+            rescale_factor,
+        )
+
+        # Init variables
+        self.constraints = constraints
+
+    def check_constraint(self, data: dict[str, Tensor]):
+        """Evaluate whether all sub-constraints are satisfied.
+
+        Args:
+            data: Model predictions and associated batch/context information.
+
+        Returns:
+            tuple[Tensor, Tensor]: A tuple `(total_satisfaction, mask)` where:
+                * `total_satisfaction`: A boolean or numeric tensor indicating
+                elementwise whether all constraints are satisfied
+                (logical AND).
+                * `mask`: A tensor of ones with the same shape as
+                `total_satisfaction`. Typically used as a weighting mask
+                in downstream processing.
         """
+        total_satisfaction: Tensor = None
+        total_mask: Tensor = None
+
+        # TODO vectorize this loop
+        for constraint in self.constraints:
+            satisfaction, mask = constraint.check_constraint(data)
+            if total_satisfaction is None:
+                total_satisfaction = satisfaction
+                total_mask = mask
+            else:
+                total_satisfaction = logical_and(total_satisfaction, satisfaction)
+                total_mask = logical_or(total_mask, mask)
 
-        # Type checking
-        validate_type("a", a, (str, Transformation))
-        validate_type("b", b, (str, Transformation))
-        validate_type("rtol", rtol, float)
-        validate_type("atol", atol, float)
-
-        # If transformation is provided, get neuron name,
-        # else use IdentityTransformation
-        if isinstance(a, Transformation):
-            neuron_name_a = a.neuron_name
-            transformation_a = a
-        else:
-            neuron_name_a = a
-            transformation_a = IdentityTransformation(neuron_name_a)
+        return total_satisfaction.float(), total_mask.float()
 
-        if isinstance(b, Transformation):
-            neuron_name_b = b.neuron_name
-            transformation_b = b
-        else:
-            neuron_name_b = b
-            transformation_b = IdentityTransformation(neuron_name_b)
+    def calculate_direction(self, data: dict[str, Tensor]):
+        """Compute the corrective direction by aggregating sub-constraint directions.
+
+        Each sub-constraint contributes its corrective direction, weighted
+        by its satisfaction mask. The directions are summed across constraints
+        for each affected layer.
+
+        Args:
+            data: Model predictions and associated batch/context information.
+
+        Returns:
+            dict[str, Tensor]: A mapping from layer identifiers to correction
+            tensors. Each entry represents the aggregated correction to apply
+            to that layer, based on the satisfaction-weighted sum of
+            sub-constraint directions.
+        """
+        total_direction: dict[str, Tensor] = {}
+
+        # TODO vectorize this loop
+        for constraint in self.constraints:
+            # TODO improve efficiency by avoiding double computation?
+            satisfaction, _ = constraint.check_constraint(data)
+            direction = constraint.calculate_direction(data)
+
+            for layer, dir in direction.items():
+                if layer not in total_direction:
+                    total_direction[layer] = satisfaction.unsqueeze(1) * dir
+                else:
+                    total_direction[layer] += satisfaction.unsqueeze(1) * dir
+
+        return total_direction
+
+
+class ORConstraint(Constraint):
+    """A composite constraint that enforces the logical OR of multiple constraints.
+
+    This class combines multiple sub-constraints and evaluates them jointly:
+
+    * The satisfaction of the OR constraint is `True` if at least one sub-constraint
+    is satisfied (elementwise logical OR).
+    * The corrective direction is computed by weighting each sub-constraint's
+    direction with its satisfaction mask and summing across all sub-constraints.
+    """
+
+    def __init__(
+        self,
+        *constraints: Constraint,
+        name: str = None,
+        monitor_only: bool = False,
+        rescale_factor: Number = 1.5,
+    ) -> None:
+        """A composite constraint that enforces the logical OR of multiple constraints.
+
+        This class combines multiple sub-constraints and evaluates them jointly:
+
+        * The satisfaction of the OR constraint is `True` if at least one sub-constraint
+        is satisfied (elementwise logical OR).
+        * The corrective direction is computed by weighting each sub-constraint's
+        direction with its satisfaction mask and summing across all sub-constraints.
+
+        Args:
+            *constraints (Constraint): One or more `Constraint` instances to be combined.
+            name (str, optional): A custom name for this constraint. If not provided,
+                the name will be composed from the sub-constraint names joined with
+                " OR ".
+            monitor_only (bool, optional): If True, the constraint will be monitored
+                but not enforced. Defaults to False.
+            rescale_factor (Number, optional): A scaling factor applied when rescaling
+                corrections. Defaults to 1.5.
+
+        Attributes:
+            constraints (tuple[Constraint, ...]): The sub-constraints being combined.
+            neurons (set): The union of neurons referenced by the sub-constraints.
+            name (str): The name of the constraint (composed or custom).
+        """
+        # Type checking
+        validate_iterable("constraints", constraints, Constraint)
 
         # Compose constraint name
-        name = f"{neuron_name_a}² + {neuron_name_b}² ≈ 1"
+        if not name:
+            name = " OR ".join([constraint.name for constraint in constraints])
 
         # Init parent class
         super().__init__(
-            {neuron_name_a, neuron_name_b},
+            set().union(*(constraint.tags for constraint in constraints)),
             name,
             monitor_only,
             rescale_factor,
         )
 
         # Init variables
-        self.transformation_a = transformation_a
-        self.transformation_b = transformation_b
-        self.rtol = rtol
-        self.atol = atol
+        self.constraints = constraints
 
-        # Get layer name and feature index from neuron_name
-        self.layer_a = self.descriptor.neuron_to_layer[neuron_name_a]
-        self.layer_b = self.descriptor.neuron_to_layer[neuron_name_b]
-        self.index_a = self.descriptor.neuron_to_index[neuron_name_a]
-        self.index_b = self.descriptor.neuron_to_index[neuron_name_b]
+    def check_constraint(self, data: dict[str, Tensor]):
+        """Evaluate whether any sub-constraints are satisfied.
 
-    def check_constraint(
-        self, prediction: dict[str, Tensor]
-    ) -> tuple[Tensor, int]:
+        Args:
+            data: Model predictions and associated batch/context information.
 
-        # Select relevant columns
-        selection_a = prediction[self.layer_a][:, self.index_a]
-        selection_b = prediction[self.layer_b][:, self.index_b]
+        Returns:
+            tuple[Tensor, Tensor]: A tuple `(total_satisfaction, mask)` where:
+                * `total_satisfaction`: A boolean or numeric tensor indicating
+                elementwise whether any constraints are satisfied
+                (logical OR).
+                * `mask`: A tensor of ones with the same shape as
+                `total_satisfaction`. Typically used as a weighting mask
+                in downstream processing.
+        """
+        total_satisfaction: Tensor = None
+        total_mask: Tensor = None
+
+        # TODO vectorize this loop
+        for constraint in self.constraints:
+            satisfaction, mask = constraint.check_constraint(data)
+            if total_satisfaction is None:
+                total_satisfaction = satisfaction
+                total_mask = mask
+            else:
+                total_satisfaction = logical_or(total_satisfaction, satisfaction)
+                total_mask = logical_or(total_mask, mask)
 
-        # Apply transformations
-        selection_a = self.transformation_a(selection_a)
-        selection_b = self.transformation_b(selection_b)
-
-        # Calculate result
-        result = isclose(
-            square(selection_a) + square(selection_b),
-            ones_like(selection_a, device=self.device),
-            rtol=self.rtol,
-            atol=self.atol,
-        ).float()
-
-        return result, numel(result)
-
-    def calculate_direction(
-        self, prediction: dict[str, Tensor]
-    ) -> Dict[str, Tensor]:
-        # NOTE currently only works for dense layers due
-        # to neuron to index translation
+        return total_satisfaction.float(), total_mask.float()
 
-        output = {}
+    def calculate_direction(self, data: dict[str, Tensor]):
+        """Compute the corrective direction by aggregating sub-constraint directions.
 
-        for layer in self.layers:
-            output[layer] = zeros_like(prediction[layer], device=self.device)
+        Each sub-constraint contributes its corrective direction, weighted
+        by its satisfaction mask. The directions are summed across constraints
+        for each affected layer.
+
+        Args:
+            data: Model predictions and associated batch/context information.
 
-        a = prediction[self.layer_a][:, self.index_a]
-        b = prediction[self.layer_b][:, self.index_b]
-        m = sqrt(square(a) + square(b))
+        Returns:
+            dict[str, Tensor]: A mapping from layer identifiers to correction
+            tensors. Each entry represents the aggregated correction to apply
+            to that layer, based on the satisfaction-weighted sum of
+            sub-constraint directions.
+        """
+        total_direction: dict[str, Tensor] = {}
 
-        output[self.layer_a][:, self.index_a] = a / m * sign(1 - m)
-        output[self.layer_b][:, self.index_b] = b / m * sign(1 - m)
+        # TODO vectorize this loop
+        for constraint in self.constraints:
+            # TODO improve efficiency by avoiding double computation?
+            satisfaction, _ = constraint.check_constraint(data)
+            direction = constraint.calculate_direction(data)
 
-        return output
+            for layer, dir in direction.items():
+                if layer not in total_direction:
+                    total_direction[layer] = satisfaction.unsqueeze(1) * dir
+                else:
+                    total_direction[layer] += satisfaction.unsqueeze(1) * dir
+
+        return total_direction