add auxiliary field to dlx.Model

pfackeldey · Nov 24, 2023 · e2aa07a · e2aa07a
1 parent b0c2e70
commit e2aa07a
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/src/dilax/model.py b/src/dilax/model.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import equinox as eqx
 import jax
@@ -96,10 +96,19 @@ def eval(model) -> jax.Array:
 
     processes: dict
     parameters: dict[str, Parameter]
+    auxiliary: Any
 
-    def __init__(self, processes: dict, parameters: dict[str, Parameter]) -> None:
+    def __init__(
+        self,
+        processes: dict,
+        parameters: dict[str, Parameter],
+        auxiliary: Any | Sentinel = _NoValue,
+    ) -> None:
         self.processes = processes
         self.parameters = parameters
+        if auxiliary is _NoValue:
+            auxiliary = {}
+        self.auxiliary = auxiliary
 
     @property
     def parameter_values(self) -> dict[str, jax.Array]: