@tree_math.vector_to_pytree_fun
def step_fn(u0):
    u = ([None] * num_steps)
    k = ([None] * num_steps)
    u[0] = u0
    k[0] = F(u0)
    for i in range(1, num_steps):
        u_star = (u0 + (dt * sum(((a[(i - 1)][j] * k[j]) for j in range(i) if a[(i - 1)][j]))))
        u[i] = P(u_star)
        k[i] = F(u[i])
    u_star = (u0 + (dt * sum(((b[j] * k[j]) for j in range(num_steps) if b[j]))))
    u_final = P(u_star)
    return u_final
@tree_math.vector_to_pytree_fun
def step_fn(u0):
    u = ([None] * num_steps)
    k = ([None] * num_steps)
    u[0] = u0
    k[0] = F(u0)
    for i in range(1, num_steps):
        u_star = (u0 + (dt * sum(((a[(i - 1)][j] * k[j]) for j in range(i) if a[(i - 1)][j]))))
        u[i] = P(u_star)
        k[i] = F(u[i])
    u_star = (u0 + (dt * sum(((b[j] * k[j]) for j in range(num_steps) if b[j]))))
    u_final = P(u_star)
    return u_final