From e2aa07abc82a2b7180f0b5a92c89d4e9d0067926 Mon Sep 17 00:00:00 2001
From: Peter Fackeldey <peter.fackeldey@rwth-aachen.de>
Date: Fri, 24 Nov 2023 13:03:23 +0100
Subject: [PATCH] add auxiliary field to dlx.Model

---
 src/dilax/model.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/dilax/model.py b/src/dilax/model.py
index 37f114c..69e32bc 100644
--- 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]: