From de53d200eb1b64b051d621546b2b0b98176f4fa9 Mon Sep 17 00:00:00 2001
From: Faezeh Habibi <155960330+Faezehabibi@users.noreply.github.com>
Date: Mon, 2 Dec 2024 14:25:58 -0500
Subject: [PATCH] Update ode_utils.py (#79)

refactor: delete @partial(jit, static_argnums=(2, )) lines

Co-authored-by: Will Gebhardt <will@gebhardts.net>
---
 ngclearn/utils/diffeq/ode_utils.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/ngclearn/utils/diffeq/ode_utils.py b/ngclearn/utils/diffeq/ode_utils.py
index d4dfb8441..fae6b08a2 100755
--- a/ngclearn/utils/diffeq/ode_utils.py
+++ b/ngclearn/utils/diffeq/ode_utils.py
@@ -55,7 +55,8 @@ def _step_forward(t, x, dx_dt, dt, x_scale): ## internal step co-routine
     _x = x * x_scale + dx_dt * dt
     return _t, _x
 
-@partial(jit, static_argnums=(2, ))
+
+
 def step_euler(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the Euler method, i.e., a
@@ -83,7 +84,7 @@ def step_euler(t, x, dfx, dt, params, x_scale=1.):
     _t, _x = _step_forward(t, x, dx_dt, dt, x_scale)
     return _t, _x
 
-@partial(jit, static_argnums=(2, ))
+
 def step_heun(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via Heun's method, i.e., a
@@ -124,7 +125,7 @@ def step_heun(t, x, dfx, dt, params, x_scale=1.):
     _, _x = _step_forward(t, x, summed_dx_dt, dt * 0.5, x_scale)
     return _t, _x
 
-@partial(jit, static_argnums=(2, ))
+
 def step_rk2(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the midpoint method, i.e., a
@@ -165,7 +166,7 @@ def step_rk2(t, x, dfx, dt, params, x_scale=1.):
 
 
 
-@partial(jit, static_argnums=(2, ))
+
 def step_rk4(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the midpoint method, i.e., a
@@ -211,7 +212,7 @@ def step_rk4(t, x, dfx, dt, params, x_scale=1.):
     _t, _x = _step_forward(t, x, _dx_dt / 6, dt, x_scale)
     return _t, _x
 
-@partial(jit, static_argnums=(2, ))
+
 def step_ralston(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via Ralston's method, i.e., a