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

congrads/__init__.py

@@ -1,21 +1,22 @@
 # __init__.py
+version = "0.2.0"
 
 # Only expose the submodules, not individual classes
-from . import core
 from . import constraints
+from . import core
 from . import datasets
 from . import descriptor
-from . import learners
 from . import metrics
 from . import networks
+from . import utils
 
 # Define __all__ to specify that the submodules are accessible, but not classes directly.
 __all__ = [
-    "core",
     "constraints",
+    "core",
     "datasets",
     "descriptor",
-    "learners",
     "metrics",
-    "networks"
-]
+    "networks",
+    "utils",
+]

congrads/constraints.py

@@ -3,7 +3,18 @@ from numbers import Number
 import random
 import string
 from typing import Callable, Dict
-from torch import Tensor, ge, gt, lt, le, zeros, FloatTensor, ones, tensor, float32
+from torch import (
+    Tensor,
+    ge,
+    gt,
+    lt,
+    le,
+    reshape,
+    stack,
+    ones,
+    tensor,
+    zeros_like,
+)
 import logging
 from torch.nn.functional import normalize
 
@@ -11,154 +22,84 @@ from .descriptor import Descriptor
 
 
 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.
-
-    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,
+        neurons: set[str],
+        name: str = None,
         rescale_factor: float = 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.
-        """
 
         # Init parent class
         super().__init__()
 
         # Init object variables
+        self.neurons = neurons
         self.rescale_factor = rescale_factor
-        self.neuron_names = neuron_names
 
         # 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."
+                f"Rescale factor for constraint {name} is <= 1. The network will favor general loss over the constraint-adjusted loss. Is this intended behaviour? Normally, the loss should always be larger than 1."
             )
 
         # 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 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}."
-            )
-
-        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."
-            )
+            self.name = f"{self.__class__.__name__}_{random_suffix}"
+            logging.warning(f"Name for constraint is not set. Using {self.name}.")
 
+        # If rescale factor is not larger than 1, warn user and adjust
         if not 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}."
+                f"Rescale factor for constraint {name} is < 1, adjusted value {rescale_factor} to {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", ...).'
                 )
 
-    @abstractmethod
-    def check_constraint(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
-        """
-        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.
-
-        Args:
-            prediction (dict[str, Tensor]): A dictionary of model predictions, indexed by layer names.
+            self.layers.add(self.descriptor.neuron_to_layer[neuron])
 
-        Returns:
-            dict[str, Tensor]: A dictionary containing the satisfaction status of the constraint for each layer or neuron.
+    # TODO only denormalize if required for efficiency
+    def _denormalize(self, input: Tensor, neuron_names: list[str]):
+        # Extract min and max for each neuron
+        min_values = tensor(
+            [self.descriptor.neuron_to_minmax[name][0] for name in neuron_names],
+            device=input.device,
+        )
+        max_values = tensor(
+            [self.descriptor.neuron_to_minmax[name][1] for name in neuron_names],
+            device=input.device,
+        )
 
-        Raises:
-            NotImplementedError: If the method is not implemented in a subclass.
-        """
+        # Apply vectorized denormalization
+        return input * (max_values - min_values) + min_values
 
+    @abstractmethod
+    def check_constraint(self, prediction: dict[str, Tensor]) -> Tensor:
         raise NotImplementedError
 
     @abstractmethod
     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.
-
-        Args:
-            prediction (dict[str, Tensor]): A dictionary of model predictions, indexed by layer names.
-
-        Returns:
-            dict[str, Tensor]: A dictionary containing the direction for each layer or neuron, to adjust the model's predictions.
-
-        Raises:
-            NotImplementedError: If the method is not implemented in a subclass.
-        """
         raise NotImplementedError
 
 
 class ScalarConstraint(Constraint):
-    """
-    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.
-
-    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,
@@ -166,23 +107,8 @@ class ScalarConstraint(Constraint):
         comparator: Callable[[Tensor, Number], Tensor],
         scalar: Number,
         name: str = None,
-        descriptor: Descriptor = None,
         rescale_factor: float = 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).
-        """
 
         # Compose constraint name
         name = f"{neuron_name}_{comparator.__name__}_{str(scalar)}"
@@ -194,10 +120,6 @@ class ScalarConstraint(Constraint):
         self.comparator = comparator
         self.scalar = scalar
 
-        if descriptor != None:
-            self.descriptor = descriptor
-            self.run_init_descriptor()
-
         # 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]
@@ -210,69 +132,31 @@ class ScalarConstraint(Constraint):
 
         # 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.
-
-        Args:
-            prediction (dict[str, Tensor]): A dictionary of model predictions, indexed by layer names.
-
-        Returns:
-            dict[str, Tensor]: A dictionary containing the constraint satisfaction result for the specified layer.
-        """
-
-        result = ~self.comparator(prediction[self.layer][:, self.index], self.scalar)
+    def check_constraint(self, prediction: dict[str, Tensor]) -> Tensor:
 
-        return {self.layer: result}
+        return ~self.comparator(prediction[self.layer][:, self.index], self.scalar)
 
     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.
+        # NOTE currently only works for dense layers due to neuron to index translation
 
-        The direction is determined by the comparator and represents either a positive or negative adjustment.
+        output = {}
 
-        Args:
-            prediction (dict[str, Tensor]): A dictionary of model predictions, indexed by layer names.
+        for layer in self.layers:
+            output[layer] = zeros_like(prediction[layer][0])
 
-        Returns:
-            dict[str, Tensor]: A dictionary containing the direction for each layer or neuron, to adjust the model's predictions.
-        """
+        output[self.layer][self.index] = self.direction
 
-        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.
-    """
 
     def __init__(
         self,
@@ -280,34 +164,22 @@ class BinaryConstraint(Constraint):
         comparator: Callable[[Tensor, Number], Tensor],
         neuron_name_right: str,
         name: str = None,
-        descriptor: Descriptor = None,
         rescale_factor: float = 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.
-        """
 
         # Compose constraint name
         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,
+            rescale_factor,
+        )
 
         # Init variables
         self.comparator = comparator
 
-        if descriptor != None:
-            self.descriptor = descriptor
-            self.run_init_descriptor()
-
         # 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]
@@ -328,65 +200,30 @@ 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]) -> Dict[str, Tensor]:
-        """
-        Checks whether the binary constraint is satisfied between the two neurons.
-
-        This function applies the comparator to the output values of the two neurons
-        and returns a Boolean result for each neuron.
-
-        Args:
-            prediction (dict[str, Tensor]): A dictionary containing the predictions for each layer.
-
-        Returns:
-            dict[str, Tensor]: A dictionary with the layer names as keys and the constraint satisfaction results as values.
-        """
+    def check_constraint(self, prediction: dict[str, Tensor]) -> Tensor:
 
-        result = ~self.comparator(
+        return ~self.comparator(
             prediction[self.layer_left][:, self.index_left],
             prediction[self.layer_right][:, self.index_right],
         )
 
-        return {self.layer_left: result, self.layer_right: result}
-
     def calculate_direction(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
-        """
-        Calculates the direction for gradient adjustment for both neurons involved in the constraint.
+        # NOTE currently only works for dense layers due to neuron to index translation
 
-        The directions are normalized and represent the direction in which the constraint should be enforced.
+        output = {}
 
-        Args:
-            prediction (dict[str, Tensor]): A dictionary containing the predictions for each layer.
+        for layer in self.layers:
+            output[layer] = zeros_like(prediction[layer][0])
 
-        Returns:
-            dict[str, Tensor]: A dictionary with the layer names as keys and the direction vectors as values.
-        """
+        output[self.layer_left][self.index_left] = self.direction_left
+        output[self.layer_right][self.index_right] = self.direction_right
 
-        output_left = zeros(
-            prediction[self.layer_left].size(),
-            device=prediction[self.layer_left].device,
-        )
-        output_left[:, self.index_left] = self.direction_left
+        for layer in self.layers:
+            output[layer] = normalize(reshape(output[layer], [1, -1]), dim=1)
 
-        output_right = zeros(
-            prediction[self.layer_right].size(),
-            device=prediction[self.layer_right].device,
-        )
-        output_right[:, self.index_right] = self.direction_right
+        return output
 
-        return {self.layer_left: output_left, self.layer_right: output_right}
 
-
-# FIXME
 class SumConstraint(Constraint):
     def __init__(
         self,
@@ -396,34 +233,17 @@ class SumConstraint(Constraint):
         weights_left: list[float] = None,
         weights_right: list[float] = None,
         name: str = None,
-        descriptor: Descriptor = None,
         rescale_factor: float = 1.5,
     ) -> None:
 
         # Init parent class
-        super().__init__(
-            set(neuron_names_left) & set(neuron_names_right), name, rescale_factor
-        )
+        neuron_names = set(neuron_names_left) | set(neuron_names_right)
+        super().__init__(neuron_names, name, rescale_factor)
 
         # Init variables
         self.comparator = comparator
-
-        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])
+        self.neuron_names_left = neuron_names_left
+        self.neuron_names_right = neuron_names_right
 
         # If comparator function is not supported, raise error
         if comparator not in [ge, le, gt, lt]:
@@ -443,13 +263,13 @@ class SumConstraint(Constraint):
 
         # 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 +279,111 @@ 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
-        )
-        self.direction_left = normalized_directions[0]
-        self.direction_right = normalized_directions[1]
+    def check_constraint(self, prediction: dict[str, Tensor]) -> Tensor:
+
+        def compute_weighted_sum(neuron_names: list[str], 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
+            ]
+
+            # Extract predictions for all neurons and apply weights in bulk
+            predictions = stack(
+                [prediction[layer][:, index] for layer, index in zip(layers, indices)],
+                dim=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)
+            # Denormalize if required
+            predictions_denorm = self._denormalize(predictions, neuron_names)
+
+            # Calculate weighted sum
+            weighted_sum = (predictions_denorm * weights.unsqueeze(0)).sum(dim=1)
 
-        # result = ~self.comparator(weighted_sum_left, weighted_sum_right)
+            return weighted_sum
 
-        # return {layer_left: result, layer_right: result}
-        pass
+        weighted_sum_left = compute_weighted_sum(
+            self.neuron_names_left, self.weights_left
+        )
+        weighted_sum_right = compute_weighted_sum(
+            self.neuron_names_right, self.weights_right
+        )
+
+        # Apply the comparator and calculate the result
+        return ~self.comparator(weighted_sum_left, weighted_sum_right)
 
     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
+        # 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])
+
+        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 MonotonicityConstraint(Constraint):
+#     # TODO docstring
+
+#     def __init__(
+#         self,
+#         neuron_name: str,
+#         name: str = None,
+#         descriptor: Descriptor = None,
+#         rescale_factor: float = 1.5,
+#     ) -> None:
+
+#         # Compose constraint name
+#         name = f"Monotonicity_{neuron_name}"
+
+#         # Init parent class
+#         super().__init__({neuron_name}, name, rescale_factor)
+
+#         # Init variables
+#         if descriptor != None:
+#             self.descriptor = descriptor
+#             self.run_init_descriptor()
+
+#         # 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]
+
+#     def check_constraint(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
+#         # Check if values for column in batch are only increasing
+#         result = ~ge(
+#             diff(
+#                 prediction[self.layer][:, self.index],
+#                 prepend=zeros_like(
+#                     prediction[self.layer][:, self.index][:1],
+#                     device=prediction[self.layer].device,
+#                 ),
+#             ),
+#             0,
+#         )
+
+#         return {self.layer: result}
+
+#     def calculate_direction(self, prediction: dict[str, Tensor]) -> Dict[str, Tensor]:
+#         # TODO implement
+
+#         output = {self.layer: zeros_like(prediction[self.layer][0])}
+#         output[self.layer][self.index] = 1
+
+#         return output