fix BH for new main function syntax

gboehl · Oct 29, 2024 · 616b329 · 616b329
1 parent 1aa0b34
commit 616b329
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/econpizza/solvers/shooting.py b/econpizza/solvers/shooting.py
@@ -101,8 +101,7 @@ def solve_current(pars, shock, XLag, XLastGuess, XPrime):
         """Solves for one period.
         """
 
-        # partial_func = jax.tree_util.Partial(func, XLag=XLag, XPrime=XPrime, XSS=stst, shocks=shock, pars=pars)
-        def partial_func(x): return func(XLag, x, XPrime, stst, shock, pars)
+        def partial_func(x): return func(XLag, x, XPrime, stst, pars, shock)
         jav = val_and_jacfwd(partial_func)
         partial_jav = jax.tree_util.Partial(jav)
         res = newton_jax_jit(partial_jav, XLastGuess, verbose=False)