Merge branch 'Konsti_Measurements' into Konsti_Recorders

examples/measurements.ipynb

@@ -45,7 +45,7 @@
     "\n",
     "loss = Loss(\n",
     "    name='loss', # Name of the measurement\n",
-    "    loss_fn=loss_fn # The function that will be called to compute the loss\n",
+    "    apply_fn=loss_fn # The function that will be called to compute the loss\n",
     ")\n",
     "print(f\"Neural state keys: {loss.neural_state_keys}\")\n",
     "\n",
@@ -140,7 +140,7 @@
     "\n",
     "loss = Loss(\n",
     "    name='loss', # Name of the measurement\n",
-    "    loss_fn=loss_fn # The function that will be called to compute the loss\n",
+    "    apply_fn=loss_fn # The function that will be called to compute the loss\n",
     ")\n",
     "\n",
     "# Defining the neural state\n",

papyrus/measurements/measurements.py

@@ -60,7 +60,7 @@ def __init__(
         name: str = "loss",
         rank: int = 0,
         public: bool = False,
-        loss_fn: Optional[Callable] = None,
+        apply_fn: Optional[Callable] = None,
     ):
         """
         Constructor method of the Loss class.
@@ -76,21 +76,22 @@ def __init__(
         public : bool (default=False)
                 Boolean flag to indicate whether the measurement resutls will be
                 accessible via a public attribute of the recorder.
-        loss_fn : Optional[Callable] (default=None)
+        apply_fn : Optional[Callable] (default=None)
                 The loss function to be used to compute the loss of the neural network.
                 If the loss function is not provided, the apply method will assume that
                 the loss is used as the input.
                 If the loss function is provided, the apply method will assume that the
                 neural network outputs and the target values are used as inputs.
         """
         super().__init__(name, rank, public)
-        self.loss_fn = loss_fn
+
+        self.apply_fn = apply_fn
 
         # Based on the provided loss function, set the apply method
-        if self.loss_fn is None:
+        if self.apply_fn is None:
             self.apply = self._apply_no_computation
         else:
-            self.apply = self.apply_computation
+            self.apply = self._apply_computation
 
         self.neural_state_keys = self._get_apply_signature()
 
@@ -137,7 +138,7 @@ def _apply_computation(
         loss : float
             The loss of the neural network.
         """
-        return self.loss_fn(predictions, targets)
+        return self.apply_fn(predictions, targets)
 
 
 class Accuracy(BaseMeasurement):
@@ -162,7 +163,7 @@ def __init__(
         name: str = "accuracy",
         rank: int = 0,
         public: bool = False,
-        accuracy_fn: Optional[Callable] = None,
+        apply_fn: Optional[Callable] = None,
     ):
         """
         Constructor method of the Accuracy class.
@@ -178,7 +179,7 @@ def __init__(
         public : bool (default=False)
                 Boolean flag to indicate whether the measurement resutls will be
                 accessible via a public attribute of the recorder.
-        accuracy_fn : Optional[Callable] (default=None)
+        apply_fn : Optional[Callable] (default=None)
                 The accuracy function to be used to compute the accuracy of the neural
                 network.
                 # If the accuracy function is not provided, the apply method will assume
@@ -187,13 +188,14 @@ def __init__(
                 the neural network outputs and the target values are used as inputs.
         """
         super().__init__(name, rank, public)
-        self.accuracy_fn = accuracy_fn
+
+        self.apply_fn = apply_fn
 
         # Based on the provided accuracy function, set the apply method
-        if self.accuracy_fn is None:
+        if self.apply_fn is None:
             self.apply = self._apply_no_computation
         else:
-            self.apply = self.apply_computation
+            self.apply = self._apply_computation
 
         self.neural_state_keys = self._get_apply_signature()
 
@@ -215,7 +217,7 @@ def _apply_no_computation(self, accuracy: Optional[float] = None) -> float:
         """
         return accuracy
 
-    def apply_computation(
+    def _apply_computation(
         self,
         predictions: Optional[np.ndarray] = None,
         targets: Optional[np.ndarray] = None,
@@ -239,7 +241,7 @@ def apply_computation(
         accuracy : float
             The accuracy of the neural network.
         """
-        return self.accuracy_fn(predictions, targets)
+        return self.apply_fn(predictions, targets)
 
 
 class NTKTrace(BaseMeasurement):