congrads 0.1.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

congrads/constraints.py

@@ -1,312 +1,512 @@
-from abc import ABC, abstractmethod
-from numbers import Number
+"""
+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.
+
+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.
+
+Dependencies:
+    - PyTorch (`torch`)
+"""
+
 import random
 import string
-from typing import Callable, Dict
-from torch import Tensor, ge, gt, lt, le, zeros, FloatTensor, ones, tensor, float32
-import logging
+import warnings
+from abc import ABC, abstractmethod
+from numbers import Number
+from typing import Callable, Dict, Union
+
+from torch import (
+    Tensor,
+    count_nonzero,
+    ge,
+    gt,
+    isclose,
+    le,
+    logical_not,
+    logical_or,
+    lt,
+    numel,
+    ones,
+    ones_like,
+    reshape,
+    sign,
+    sqrt,
+    square,
+    stack,
+    tensor,
+    zeros_like,
+)
 from torch.nn.functional import normalize
 
 from .descriptor import Descriptor
+from .transformations import IdentityTransformation, Transformation
+from .utils import validate_comparator_pytorch, validate_iterable, validate_type
 
 
 class Constraint(ABC):
     """
-    Abstract base class for defining constraints that can be applied during optimization in the constraint-guided gradient descent process.
-
-    A constraint guides the optimization by evaluating the model's predictions and adjusting the loss based on certain conditions.
-    Constraints can be applied to specific layers or neurons of the model, and they are scaled by a rescale factor to control the influence of the constraint on the overall loss.
+    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
+    during training and can adjust loss to enforce constraints. Subclasses
+    must implement the `check_constraint` and `calculate_direction` methods.
+
+    Args:
+        neurons (set[str]): Names of the neurons 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.
+        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
+            defined in the `descriptor`.
+
+    Note:
+        - If `rescale_factor <= 1`, a warning is issued, and the value is
+          adjusted to a positive value greater than 1.
+        - If `name` is not provided, a name is auto-generated,
+          and a warning is logged.
 
-    Attributes:
-        descriptor (Descriptor): The descriptor object that provides a mapping of neurons to layers.
-        constraint_name (str): A unique name for the constraint, which can be provided or generated automatically.
-        rescale_factor (float): A factor used to scale the influence of the constraint on the overall loss.
-        neuron_names (set[str]): A set of neuron names that are involved in the constraint.
-        layers (set): A set of layers associated with the neurons specified in `neuron_names`.
     """
 
     descriptor: Descriptor = None
+    device = None
 
     def __init__(
         self,
-        neuron_names: set[str],
-        constraint_name: str = None,
-        rescale_factor: float = 1.5,
+        neurons: set[str],
+        name: str = None,
+        monitor_only: bool = False,
+        rescale_factor: Number = 1.5,
     ) -> None:
         """
-        Initializes the Constraint object with the given neuron names, constraint name, and rescale factor.
-
-        Args:
-            neuron_names (set[str]): A set of neuron names that are affected by the constraint.
-            constraint_name (str, optional): A custom name for the constraint. If not provided, a random name is generated.
-            rescale_factor (float, optional): A factor that scales the influence of the constraint. Defaults to 1.5.
-
-        Raises:
-            ValueError: If the descriptor has not been set or if a neuron name is not found in the descriptor.
+        Initializes a new Constraint instance.
         """
 
         # Init parent class
         super().__init__()
 
+        # Type checking
+        validate_iterable("neurons", neurons, str)
+        validate_type("name", name, (str, type(None)))
+        validate_type("monitor_only", monitor_only, bool)
+        validate_type("rescale_factor", rescale_factor, Number)
+
         # Init object variables
+        self.neurons = neurons
         self.rescale_factor = rescale_factor
-        self.neuron_names = neuron_names
+        self.monitor_only = monitor_only
 
         # Perform checks
         if rescale_factor <= 1:
-            logging.warning(
-                f"Rescale factor for constraint {constraint_name} is <= 1. The network will favour general loss over the constraint-adjusted loss. Is this intended behaviour? Normally, the loss should always be larger than 1."
+            warnings.warn(
+                "Rescale factor for constraint %s is <= 1. The network \
+                    will favor general loss over the constraint-adjusted loss. \
+                    Is this intended behavior? Normally, the loss should \
+                    always be larger than 1.",
+                name,
             )
 
-        # If no constraint_name is set, generate one based on the class name and a random suffix
-        if constraint_name:
-            self.constraint_name = constraint_name
+        # 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)
             )
-            self.constraint_name = f"{self.__class__.__name__}_{random_suffix}"
-            logging.warning(
-                f"Name for constraint is not set. Using {self.constraint_name}."
+            self.name = f"{self.__class__.__name__}_{random_suffix}"
+            warnings.warn(
+                "Name for constraint is not set. Using %s.", self.name
             )
 
-        if self.descriptor == None:
-            raise ValueError(
-                "The descriptor of the base Constraint class in not set. Please assign the descriptor to the general Constraint class with 'Constraint.descriptor = descriptor' before defining network-specific contraints."
-            )
-
-        if not rescale_factor > 1:
+        # If rescale factor is not larger than 1, warn user and adjust
+        if rescale_factor <= 1:
             self.rescale_factor = abs(rescale_factor) + 1.5
-            logging.warning(
-                f"Rescale factor for constraint {constraint_name} is < 1, adjusted value {rescale_factor} to {self.rescale_factor}."
+            warnings.warn(
+                "Rescale factor for constraint %s is < 1, adjusted value \
+                    %s to %s.",
+                name,
+                rescale_factor,
+                self.rescale_factor,
             )
         else:
             self.rescale_factor = rescale_factor
 
-        self.neuron_names = neuron_names
-
-        self.run_init_descriptor()
-
-    def run_init_descriptor(self) -> None:
-        """
-        Initializes the layers associated with the constraint by mapping the neuron names to their corresponding layers
-        from the descriptor.
-
-        This method populates the `layers` attribute with layers associated with the neuron names provided in the constraint.
-
-        Raises:
-            ValueError: If a neuron name is not found in the descriptor's mapping of neurons to layers.
-        """
-
+        # Infer layers from descriptor and neurons
         self.layers = set()
-        for neuron_name in self.neuron_names:
-            if neuron_name in self.descriptor.neuron_to_layer.keys():
-                self.layers.add(self.descriptor.neuron_to_layer[neuron_name])
-            else:
+        for neuron in self.neurons:
+            if neuron not in self.descriptor.neuron_to_layer.keys():
                 raise ValueError(
-                    f'The neuron name {neuron_name} used with constraint {self.constraint_name} is not defined in the descriptor. Please add it to the correct layer using descriptor.add("layer", ...).'
+                    f'The neuron name {neuron} used with constraint \
+                        {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])
+
     @abstractmethod
-    def check_constraint(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
+    def check_constraint(
+        self, prediction: dict[str, Tensor]
+    ) -> tuple[Tensor, int]:
         """
-        Abstract method to check if the constraint is satisfied based on the model's predictions.
-
-        This method should be implemented in subclasses to define the specific logic for evaluating the constraint based on the model's predictions.
+        Evaluates whether the given model predictions satisfy the constraint.
 
         Args:
-            prediction (dict[str, Tensor]): A dictionary of model predictions, indexed by layer names.
+            prediction (dict[str, Tensor]): Model predictions for the neurons.
 
         Returns:
-            dict[str, Tensor]: A dictionary containing the satisfaction status of the constraint for each layer or neuron.
+            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.
 
         Raises:
-            NotImplementedError: If the method is not implemented in a subclass.
+            NotImplementedError: If not implemented in a subclass.
         """
 
         raise NotImplementedError
 
     @abstractmethod
-    def calculate_direction(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
+    def calculate_direction(
+        self, prediction: dict[str, Tensor]
+    ) -> Dict[str, Tensor]:
         """
-        Abstract method to calculate the direction in which the model's predictions need to be adjusted to satisfy the constraint.
-
-        This method should be implemented in subclasses to define how to adjust the model's predictions based on the constraint.
+        Calculates adjustment directions for neurons to
+        better satisfy the constraint.
 
         Args:
-            prediction (dict[str, Tensor]): A dictionary of model predictions, indexed by layer names.
+            prediction (dict[str, Tensor]): Model predictions for the neurons.
 
         Returns:
-            dict[str, Tensor]: A dictionary containing the direction for each layer or neuron, to adjust the model's predictions.
+            Dict[str, Tensor]: Dictionary mapping neuron layers to tensors
+            specifying the adjustment direction for each neuron.
 
         Raises:
-            NotImplementedError: If the method is not implemented in a subclass.
+            NotImplementedError: If not implemented in a subclass.
         """
+
         raise NotImplementedError
 
 
-class ScalarConstraint(Constraint):
+class ImplicationConstraint(Constraint):
+    """
+    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.
+
     """
-    A subclass of the `Constraint` class that applies a scalar constraint on a specific neuron in the model.
 
-    This constraint compares the value of a specific neuron in the model to a scalar value using a specified comparator (e.g., greater than, less than).
-    If the constraint is violated, it adjusts the loss according to the direction defined by the comparator.
+    def __init__(
+        self,
+        head: Constraint,
+        body: Constraint,
+        name=None,
+        monitor_only=False,
+        rescale_factor=1.5,
+    ):
+        """
+        Initializes an ImplicationConstraint instance.
+        """
+
+        # Type checking
+        validate_type("head", head, Constraint)
+        validate_type("body", body, Constraint)
+
+        # Compose constraint name
+        name = f"{body.name} if {head.name}"
+
+        # Init parent class
+        super().__init__(
+            head.neurons | body.neurons,
+            name,
+            monitor_only,
+            rescale_factor,
+        )
+
+        self.head = head
+        self.body = body
+
+    def check_constraint(
+        self, prediction: dict[str, Tensor]
+    ) -> tuple[Tensor, int]:
+
+        # Check satisfaction of head and body constraints
+        head_satisfaction, _ = self.head.check_constraint(prediction)
+        body_satisfaction, _ = self.body.check_constraint(prediction)
+
+        # 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()
+
+        return result, count_nonzero(head_satisfaction)
+
+    def calculate_direction(
+        self, prediction: dict[str, Tensor]
+    ) -> Dict[str, Tensor]:
+        # NOTE currently only works for dense layers
+        # due to neuron to index translation
+
+        # Use directions of constraint body as update vector
+        return self.body.calculate_direction(prediction)
+
+
+class ScalarConstraint(Constraint):
+    """
+    A constraint that enforces scalar-based comparisons on a specific neuron.
+
+    This class ensures that the output of a specified neuron 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
+            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.
+        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 `neuron_name` must be defined in the `descriptor` mapping.
+        - The constraint name is composed using the neuron name,
+          comparator, and scalar value.
 
-    Attributes:
-        comparator (Callable[[Tensor, Number], Tensor]): A comparator function (e.g., greater than, less than) to evaluate the constraint.
-        scalar (Number): The scalar value to compare the neuron value against.
-        direction (int): The direction in which the constraint should adjust the model's predictions (either 1 or -1 based on the comparator).
-        layer (str): The layer associated with the specified neuron.
-        index (int): The index of the specified neuron within the layer.
     """
 
     def __init__(
         self,
-        neuron_name: str,
+        operand: Union[str, Transformation],
         comparator: Callable[[Tensor, Number], Tensor],
         scalar: Number,
         name: str = None,
-        descriptor: Descriptor = None,
-        rescale_factor: float = 1.5,
+        monitor_only: bool = False,
+        rescale_factor: Number = 1.5,
     ) -> None:
         """
-        Initializes the ScalarConstraint with the given neuron name, comparator, scalar value, and other optional parameters.
-
-        Args:
-            neuron_name (str): The name of the neuron that the constraint applies to.
-            comparator (Callable[[Tensor, Number], Tensor]): The comparator function used to evaluate the constraint (e.g., ge, le, gt, lt).
-            scalar (Number): The scalar value that the neuron value is compared to.
-            name (str, optional): A custom name for the constraint. If not provided, a name is generated based on the neuron name, comparator, and scalar.
-            descriptor (Descriptor, optional): The descriptor that maps neurons to layers. If not provided, the global descriptor is used.
-            rescale_factor (float, optional): A factor that scales the influence of the constraint on the overall loss. Defaults to 1.5.
-
-        Raises:
-            ValueError: If the comparator function is not one of the supported comparison operators (ge, le, gt, lt).
+        Initializes a ScalarConstraint instance.
         """
 
+        # 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 isinstance(operand, Transformation):
+            neuron_name = operand.neuron_name
+            transformation = operand
+        else:
+            neuron_name = operand
+            transformation = IdentityTransformation(neuron_name)
+
         # Compose constraint name
-        name = f"{neuron_name}_{comparator.__name__}_{str(scalar)}"
+        name = f"{neuron_name} {comparator.__name__} {str(scalar)}"
 
         # Init parent class
-        super().__init__({neuron_name}, name, rescale_factor)
+        super().__init__({neuron_name}, name, monitor_only, rescale_factor)
 
         # Init variables
         self.comparator = comparator
         self.scalar = scalar
-
-        if descriptor != None:
-            self.descriptor = descriptor
-            self.run_init_descriptor()
+        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]
 
-        # If comparator function is not supported, raise error
-        if comparator not in [ge, le, gt, lt]:
-            raise ValueError(
-                f"Comparator {str(comparator)} used for constraint {name} is not supported. Only ge, le, gt, lt are allowed."
-            )
-
         # 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, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
-        """
-        Checks if the constraint is satisfied based on the model's predictions.
-
-        The constraint is evaluated by applying the comparator to the value of the specified neuron and the scalar value.
+    def check_constraint(
+        self, prediction: dict[str, Tensor]
+    ) -> tuple[Tensor, int]:
 
-        Args:
-            prediction (dict[str, Tensor]): A dictionary of model predictions, indexed by layer names.
+        # Select relevant columns
+        selection = prediction[self.layer][:, self.index]
 
-        Returns:
-            dict[str, Tensor]: A dictionary containing the constraint satisfaction result for the specified layer.
-        """
+        # Apply transformation
+        selection = self.transformation(selection)
 
-        result = ~self.comparator(prediction[self.layer][:, self.index], self.scalar)
+        # Calculate current constraint result
+        result = self.comparator(selection, self.scalar).float()
+        return result, numel(result)
 
-        return {self.layer: 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
 
-    def calculate_direction(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
-        """
-        Calculates the direction in which the model's predictions need to be adjusted to satisfy the constraint.
+        output = {}
 
-        The direction is determined by the comparator and represents either a positive or negative adjustment.
+        for layer in self.layers:
+            output[layer] = zeros_like(prediction[layer][0], device=self.device)
 
-        Args:
-            prediction (dict[str, Tensor]): A dictionary of model predictions, indexed by layer names.
+        output[self.layer][self.index] = self.direction
 
-        Returns:
-            dict[str, Tensor]: A dictionary containing the direction for each layer or neuron, to adjust the model's predictions.
-        """
-
-        output = zeros(
-            prediction[self.layer].size(),
-            device=prediction[self.layer].device,
-        )
-        output[:, self.index] = self.direction
+        for layer in self.layers:
+            output[layer] = normalize(reshape(output[layer], [1, -1]), dim=1)
 
-        return {self.layer: output}
+        return output
 
 
 class BinaryConstraint(Constraint):
     """
-    A class representing a binary constraint between two neurons in a neural network.
-
-    This class checks and enforces a constraint between two neurons using a
-    comparator function. The constraint is applied between two neurons located
-    in different layers of the neural network. The class also calculates the
-    direction for gradient adjustment based on the comparator.
-
-    Attributes:
-        neuron_name_left (str): The name of the first neuron involved in the constraint.
-        neuron_name_right (str): The name of the second neuron involved in the constraint.
-        comparator (Callable[[Tensor, Number], Tensor]): A function that compares the values of the two neurons.
-        layer_left (str): The layer name for the first neuron.
-        layer_right (str): The layer name for the second neuron.
-        index_left (int): The index of the first neuron within its layer.
-        index_right (int): The index of the second neuron within its layer.
-        direction_left (float): The normalized direction for gradient adjustment of the first neuron.
-        direction_right (float): The normalized direction for gradient adjustment of the second neuron.
+    A constraint that enforces a binary comparison between two neurons.
+
+    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.
+
+    Args:
+        operand_left (Union[str, Transformation]): Name of the left
+            neuron 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.
+        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.
+        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 neuron names must be defined in the `descriptor` mapping.
+        - The constraint name is composed using the left neuron name,
+          comparator, and right neuron name.
+
     """
 
     def __init__(
         self,
-        neuron_name_left: str,
+        operand_left: Union[str, Transformation],
         comparator: Callable[[Tensor, Number], Tensor],
-        neuron_name_right: str,
+        operand_right: Union[str, Transformation],
         name: str = None,
-        descriptor: Descriptor = None,
-        rescale_factor: float = 1.5,
+        monitor_only: bool = False,
+        rescale_factor: Number = 1.5,
     ) -> None:
         """
-        Initializes the binary constraint with two neurons, a comparator, and other configuration options.
-
-        Args:
-            neuron_name_left (str): The name of the first neuron in the constraint.
-            comparator (Callable[[Tensor, Number], Tensor]): A function that compares the values of the two neurons.
-            neuron_name_right (str): The name of the second neuron in the constraint.
-            name (str, optional): The name of the constraint. If not provided, a default name is generated.
-            descriptor (Descriptor, optional): The descriptor containing the mapping of neurons to layers.
-            rescale_factor (float, optional): A factor to rescale the constraint value. Default is 1.5.
+        Initializes a BinaryConstraint instance.
         """
 
+        # 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 isinstance(operand_left, Transformation):
+            neuron_name_left = operand_left.neuron_name
+            transformation_left = operand_left
+        else:
+            neuron_name_left = operand_left
+            transformation_left = IdentityTransformation(neuron_name_left)
+
+        if isinstance(operand_right, Transformation):
+            neuron_name_right = operand_right.neuron_name
+            transformation_right = operand_right
+        else:
+            neuron_name_right = operand_right
+            transformation_right = IdentityTransformation(neuron_name_right)
+
         # Compose constraint name
-        name = f"{neuron_name_left}_{comparator.__name__}_{neuron_name_right}"
+        name = f"{neuron_name_left} {comparator.__name__} {neuron_name_right}"
 
         # Init parent class
-        super().__init__({neuron_name_left, neuron_name_right}, name, rescale_factor)
+        super().__init__(
+            {neuron_name_left, neuron_name_right},
+            name,
+            monitor_only,
+            rescale_factor,
+        )
 
         # Init variables
         self.comparator = comparator
-
-        if descriptor != None:
-            self.descriptor = descriptor
-            self.run_init_descriptor()
+        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]
@@ -314,12 +514,6 @@ class BinaryConstraint(Constraint):
         self.index_left = self.descriptor.neuron_to_index[neuron_name_left]
         self.index_right = self.descriptor.neuron_to_index[neuron_name_right]
 
-        # If comparator function is not supported, raise error
-        if comparator not in [ge, le, gt, lt]:
-            raise RuntimeError(
-                f"Comparator {str(comparator)} used for constraint {name} is not supported. Only ge, le, gt, lt are allowed."
-            )
-
         # Calculate directions based on constraint operator
         if self.comparator in [lt, le]:
             self.direction_left = -1
@@ -328,128 +522,179 @@ class BinaryConstraint(Constraint):
             self.direction_left = 1
             self.direction_right = -1
 
-        # Normalize directions
-        normalized_directions = normalize(
-            tensor([self.direction_left, self.direction_right]).type(float32),
-            p=2,
-            dim=0,
-        )
-        self.direction_left = normalized_directions[0]
-        self.direction_right = normalized_directions[1]
+    def check_constraint(
+        self, prediction: dict[str, Tensor]
+    ) -> tuple[Tensor, int]:
 
-    def check_constraint(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
-        """
-        Checks whether the binary constraint is satisfied between the two neurons.
+        # Select relevant columns
+        selection_left = prediction[self.layer_left][:, self.index_left]
+        selection_right = prediction[self.layer_right][:, self.index_right]
 
-        This function applies the comparator to the output values of the two neurons
-        and returns a Boolean result for each neuron.
+        # Apply transformations
+        selection_left = self.transformation_left(selection_left)
+        selection_right = self.transformation_right(selection_right)
 
-        Args:
-            prediction (dict[str, Tensor]): A dictionary containing the predictions for each layer.
+        result = self.comparator(selection_left, selection_right).float()
 
-        Returns:
-            dict[str, Tensor]: A dictionary with the layer names as keys and the constraint satisfaction results as values.
-        """
+        return result, numel(result)
 
-        result = ~self.comparator(
-            prediction[self.layer_left][:, self.index_left],
-            prediction[self.layer_right][:, self.index_right],
-        )
+    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 {self.layer_left: result, self.layer_right: result}
+        output = {}
 
-    def calculate_direction(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
-        """
-        Calculates the direction for gradient adjustment for both neurons involved in the constraint.
+        for layer in self.layers:
+            output[layer] = zeros_like(prediction[layer][0], device=self.device)
 
-        The directions are normalized and represent the direction in which the constraint should be enforced.
+        output[self.layer_left][self.index_left] = self.direction_left
+        output[self.layer_right][self.index_right] = self.direction_right
 
-        Args:
-            prediction (dict[str, Tensor]): A dictionary containing the predictions for each layer.
+        for layer in self.layers:
+            output[layer] = normalize(reshape(output[layer], [1, -1]), dim=1)
 
-        Returns:
-            dict[str, Tensor]: A dictionary with the layer names as keys and the direction vectors as values.
-        """
+        return output
 
-        output_left = zeros(
-            prediction[self.layer_left].size(),
-            device=prediction[self.layer_left].device,
-        )
-        output_left[:, self.index_left] = self.direction_left
 
-        output_right = zeros(
-            prediction[self.layer_right].size(),
-            device=prediction[self.layer_right].device,
-        )
-        output_right[:, self.index_right] = self.direction_right
-
-        return {self.layer_left: output_left, self.layer_right: output_right}
+class SumConstraint(Constraint):
+    """
+    A constraint that enforces a weighted summation comparison
+    between two groups of neurons.
+
+    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.
 
+    """
 
-# FIXME
-class SumConstraint(Constraint):
     def __init__(
         self,
-        neuron_names_left: list[str],
+        operands_left: list[Union[str, Transformation]],
         comparator: Callable[[Tensor, Number], Tensor],
-        neuron_names_right: list[str],
-        weights_left: list[float] = None,
-        weights_right: list[float] = None,
+        operands_right: list[Union[str, Transformation]],
+        weights_left: list[Number] = None,
+        weights_right: list[Number] = None,
         name: str = None,
-        descriptor: Descriptor = None,
-        rescale_factor: float = 1.5,
+        monitor_only: bool = False,
+        rescale_factor: Number = 1.5,
     ) -> None:
+        """
+        Initializes the SumConstraint.
+        """
 
-        # Init parent class
-        super().__init__(
-            set(neuron_names_left) & set(neuron_names_right), name, rescale_factor
+        # 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("weights_left", weights_left, 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] = []
+        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)
+                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)
+                )
+
+        neuron_names_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)
+                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)
+                )
+
+        # 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)
+        )
+        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)
+
         # Init variables
         self.comparator = comparator
+        self.neuron_names_left = neuron_names_left
+        self.neuron_names_right = neuron_names_right
+        self.transformations_left = transformations_left
+        self.transformations_right = transformations_right
 
-        if descriptor != None:
-            self.descriptor = descriptor
-            self.run_init_descriptor()
-
-        # Get layer names and feature indices from neuron_name
-        self.layers_left = []
-        self.indices_left = []
-        for neuron_name in neuron_names_left:
-            self.layers_left.append(self.descriptor.neuron_to_layer[neuron_name])
-            self.indices_left.append(self.descriptor.neuron_to_index[neuron_name])
-
-        self.layers_right = []
-        self.indices_right = []
-        for neuron_name in neuron_names_right:
-            self.layers_right.append(self.descriptor.neuron_to_layer[neuron_name])
-            self.indices_right.append(self.descriptor.neuron_to_index[neuron_name])
-
-        # If comparator function is not supported, raise error
-        if comparator not in [ge, le, gt, lt]:
-            raise ValueError(
-                f"Comparator {str(comparator)} used for constraint {name} is not supported. Only ge, le, gt, lt are allowed."
-            )
-
-        # If feature list dimensions don't match weight list dimensions, raise error
+        # If feature list dimensions don't match
+        # weight list dimensions, raise error
         if weights_left and (len(neuron_names_left) != len(weights_left)):
             raise ValueError(
-                "The dimensions of neuron_names_left don't match with the dimensions of weights_left."
+                "The dimensions of neuron_names_left don't match with the \
+                    dimensions of weights_left."
             )
         if weights_right and (len(neuron_names_right) != len(weights_right)):
             raise ValueError(
-                "The dimensions of neuron_names_right don't match with the dimensions of weights_right."
+                "The dimensions of neuron_names_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 = FloatTensor(weights_left)
+            self.weights_left = tensor(weights_left, device=self.device)
         else:
-            self.weights_left = ones(len(neuron_names_left))
+            self.weights_left = ones(len(neuron_names_left), device=self.device)
         if weights_right:
-            self.weights_right = FloatTensor(weights_right)
+            self.weights_right = tensor(weights_right, device=self.device)
         else:
-            self.weights_right = ones(len(neuron_names_right))
+            self.weights_right = ones(
+                len(neuron_names_right), device=self.device
+            )
 
         # Calculate directions based on constraint operator
         if self.comparator in [lt, le]:
@@ -459,49 +704,216 @@ class SumConstraint(Constraint):
             self.direction_left = 1
             self.direction_right = -1
 
-        # Normalize directions
-        normalized_directions = normalize(
-            tensor(self.direction_left, self.direction_right), p=2, dim=0
+    def check_constraint(
+        self, prediction: dict[str, Tensor]
+    ) -> tuple[Tensor, int]:
+
+        def compute_weighted_sum(
+            neuron_names: 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)
+            ]
+
+            # Apply transformations
+            results = []
+            for transformation, selection in zip(transformations, selections):
+                results.append(transformation(selection))
+
+            # Extract predictions for all neurons and apply weights in bulk
+            predictions = stack(
+                results,
+                dim=1,
+            )
+
+            # Calculate weighted sum
+            return (predictions * weights.unsqueeze(0)).sum(dim=1)
+
+        # Compute weighted sums
+        weighted_sum_left = compute_weighted_sum(
+            self.neuron_names_left,
+            self.transformations_left,
+            self.weights_left,
+        )
+        weighted_sum_right = compute_weighted_sum(
+            self.neuron_names_right,
+            self.transformations_right,
+            self.weights_right,
         )
-        self.direction_left = normalized_directions[0]
-        self.direction_right = normalized_directions[1]
 
-    def check_constraint(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
-        raise NotImplementedError
-        # # TODO remove the dynamic to device conversion and do this in initialization one way or another
-        # weighted_sum_left = (
-        #     prediction[layer_left][:, index_left]
-        #     * self.weights_left.to(prediction[layer_left].device)
-        # ).sum(dim=1)
-        # weighted_sum_right = (
-        #     prediction[layer_right][:, index_right]
-        #     * self.weights_right.to(prediction[layer_right].device)
-        # ).sum(dim=1)
-
-        # result = ~self.comparator(weighted_sum_left, weighted_sum_right)
-
-        # return {layer_left: result, layer_right: result}
-        pass
-
-    def calculate_direction(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
-        raise NotImplementedError
-        # # TODO move this to constructor somehow
-        # layer_left = prediction.neuron_to_layer[self.neuron_name_left]
-        # layer_right = prediction.neuron_to_layer[self.neuron_name_right]
-        # index_left = prediction.neuron_to_index[self.neuron_name_left]
-        # index_right = prediction.neuron_to_index[self.neuron_name_right]
-
-        # output_left = zeros(
-        #     prediction[layer_left].size(),
-        #     device=prediction[layer_left].device,
-        # )
-        # output_left[:, index_left] = self.direction_left
-
-        # output_right = zeros(
-        #     prediction.layer_to_data[layer_right].size(),
-        #     device=prediction.layer_to_data[layer_right].device,
-        # )
-        # output_right[:, index_right] = self.direction_right
-
-        # return {layer_left: output_left, layer_right: output_right}
-        pass
+        # Apply the comparator and calculate the result
+        result = self.comparator(weighted_sum_left, weighted_sum_right).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
+
+        output = {}
+
+        for layer in self.layers:
+            output[layer] = zeros_like(prediction[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]
+            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]
+            output[layer][index] = self.direction_right
+
+        for layer in self.layers:
+            output[layer] = normalize(reshape(output[layer], [1, -1]), dim=1)
+
+        return output
+
+
+class PythagoreanIdentityConstraint(Constraint):
+    """
+    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.
+
+    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).
+
+    Raises:
+        TypeError: If a provided attribute has an incompatible type.
+
+    """
+
+    def __init__(
+        self,
+        a: Union[str, Transformation],
+        b: Union[str, Transformation],
+        rtol: float = 0.00001,
+        atol: float = 1e-8,
+        name: str = None,
+        monitor_only: bool = False,
+        rescale_factor: Number = 1.5,
+    ) -> None:
+        """
+        Initialize the PythagoreanIdentityConstraint.
+        """
+
+        # 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)
+
+        if isinstance(b, Transformation):
+            neuron_name_b = b.neuron_name
+            transformation_b = b
+        else:
+            neuron_name_b = b
+            transformation_b = IdentityTransformation(neuron_name_b)
+
+        # Compose constraint name
+        name = f"{neuron_name_a}² + {neuron_name_b}² ≈ 1"
+
+        # Init parent class
+        super().__init__(
+            {neuron_name_a, neuron_name_b},
+            name,
+            monitor_only,
+            rescale_factor,
+        )
+
+        # Init variables
+        self.transformation_a = transformation_a
+        self.transformation_b = transformation_b
+        self.rtol = rtol
+        self.atol = atol
+
+        # 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, prediction: dict[str, Tensor]
+    ) -> tuple[Tensor, int]:
+
+        # Select relevant columns
+        selection_a = prediction[self.layer_a][:, self.index_a]
+        selection_b = prediction[self.layer_b][:, self.index_b]
+
+        # 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
+
+        output = {}
+
+        for layer in self.layers:
+            output[layer] = zeros_like(prediction[layer], device=self.device)
+
+        a = prediction[self.layer_a][:, self.index_a]
+        b = prediction[self.layer_b][:, self.index_b]
+        m = sqrt(square(a) + square(b))
+
+        output[self.layer_a][:, self.index_a] = a / m * sign(1 - m)
+        output[self.layer_b][:, self.index_b] = b / m * sign(1 - m)
+
+        return output