Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Changes from PR1. Scrubbed notebook outputs, incorporated lax looping… #2

Merged
merged 1 commit into from
Feb 3, 2021

Conversation

simbilod
Copy link
Contributor

@simbilod simbilod commented Feb 3, 2021

… function, and cosmetic changes to notebook 2.

Thanks for the feedback!

Good calls on notebook outputs -- here's a clean history without them. On my end with GPU the lax function is faster to compile, but seems slower to execute (by a factor of ~6).

… function, and cosmetic changes to notebook 2.
@flaport
Copy link
Owner

flaport commented Feb 3, 2021

Thanks a lot Simon!

I have indeed noticed it too that using jax.lax.scan seems to be slower to execute. Maybe because the loop is no longer being unrolled? I have no idea why this is to be honest.

Anyway thanks for this PR! And... always welcome to add new stuff in the future 😉

[closes #1 ]

@flaport flaport merged commit 64e5e2f into flaport:master Feb 3, 2021
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.

2 participants