Make all JacobianHandler arguments keyword-args for now

Author: pnkraemer

probdiffeq/probdiffeq.py

@@ -172,23 +172,23 @@ def init_jacobian_handler(self):
         """
         raise NotImplementedError
 
-    def materialize_dense(self, fun, x, state):
+    def materialize_dense(self, fun, x, state, /):
         """Materialize a dense Jacobian.
 
         This is typically used for first-order linearization in dense
         state-space models.
         """
         raise NotImplementedError
 
-    def calculate_trace(self, fun, x, state):
+    def calculate_trace(self, fun, x, state, /):
         """Calculate the trace of a Jacobian.
 
         This is typically used for first-order linearization in isotropic
         state-space models.
         """
         raise NotImplementedError
 
-    def calculate_diagonal(self, fun, x, state):
+    def calculate_diagonal(self, fun, x, state, /):
         """Calculate the diagonal of a Jacobian.
 
         This is typically used for first-order linearization in block-diagonal
@@ -209,20 +209,20 @@ def __init__(self, *, jacfun=func.jacfwd):
     def init_jacobian_handler(self):
         return ()
 
-    def materialize_dense(self, fun, x, state):
+    def materialize_dense(self, fun, x, state, /):
         del state
         fx = fun(x)
         dfx = func.jacfwd(fun)(x)
         return fx, dfx, ()
 
-    def calculate_trace(self, fun, x, state):
+    def calculate_trace(self, fun, x, state, /):
         del state
         fx = fun(x)
         dfx = func.jacfwd(fun)(x)
         dfx_trace = linalg.trace(dfx)
         return fx, dfx_trace, ()
 
-    def calculate_diagonal(self, fun, x, state):
+    def calculate_diagonal(self, fun, x, state, /):
         del state
         fx = fun(x)
         dfx = func.jacfwd(fun)(x)
@@ -245,7 +245,7 @@ def __init__(self, *, seed=1, num_probes=10):
     def init_jacobian_handler(self):
         return random.prng_key(seed=self.seed)
 
-    def materialize_dense(self, fun, x, state):
+    def materialize_dense(self, fun, x, state, /):
         # TODO: approximate Jacobian with outer products instead of forming?
         # What is the "correct" thing to do?
         fx = fun(x)
@@ -288,7 +288,7 @@ def __init__(self, *, seed=1, num_probes=10):
     def init_jacobian_handler(self):
         return random.prng_key(seed=self.seed)
 
-    def materialize_dense(self, fun, x, state):
+    def materialize_dense(self, fun, x, state, /):
         # TODO: approximate Jacobian with outer products instead of forming?
         # What is the "correct" thing to do?
         fx = fun(x)