Skip to content

Conversation

@zqigolden
Copy link
Collaborator

Resolve #18

@zqigolden zqigolden requested a review from mfkasim1 June 28, 2024 00:32
Comment on lines 190 to 196
def success_fn(carry, i):
yi, _ = carry
xi = xinp[i-1]
dx = xinp[i] - xinp[i-1]
dt = tpts[i] - tpts[i-1]
yn, success = self.ode_step(func, yi, xi, dt, dx, params)
return yn, success
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like xinp and tpts variables are coming from outside this function and that would make this function unpure. JAX has a strong obsession with pure function (see this for example), so can you make it a pure function?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed with make function purer

@zqigolden zqigolden requested a review from mfkasim1 July 12, 2024 12:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Write some ODE solvers

3 participants