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.

Sign up for GitHub

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

Jump to bottom

Add a proof-of-concept for an Algorithm that uses Jax for its forward/backward passes [RT-71] #4

Merged
merged 27 commits into from
Jun 14, 2024
Merged

Add a proof-of-concept for an Algorithm that uses Jax for its forward/backward passes [RT-71] #4

merged 27 commits into from
Jun 14, 2024

Conversation

lebrice
Copy link
Collaborator

@lebrice lebrice commented May 31, 2024

No description provided.

@lebrice
Add an example algo that uses jax!
61c0662 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Simplify the jax example
752afd1 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Slightly tweak the jax example
d5b400c 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Tweak the jax example
1a71d48 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Tweak algo a bit (again)
4deb156 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Use flax nn.Module
6306ccd 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Hacky: Wrap jax fn into a torch.autograd.Function
a52972c 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Make it work with automatic optimization and jit!
8feece6 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Able to use jax in intermediate node in graph!
8d79e67 
Signed-off-by: Fabrice Normandin <[email protected]>
lebrice added 18 commits June 6, 2024 21:49
@lebrice
Update to use git packages
0473808 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Rename batch_idx->batch_index everywhere
70e7053 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Fix broken callback due to batch_idx rename
2c53544 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Use a callback to log classification metrics
e3f4c90 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Update the jax algo
18f6c25 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Make the callback compatible with more recent PL
fcd8484 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Make the Jax algo usable from CLI, tweak configs
8f595fd 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Fix tests to use the tensor_regression package
75765e9 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Fix some issues with config registration in tests
61e204f 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Fix other tiny issues in test code
9a220a1 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Fix issue with resnet50 config
542d086 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Add some generated tests for the Jax algo example
00cbd4f 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Fix tests for algo that doesnt support jax
cdd548c 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
'fix' issue with doctest of some configs
3cf5cfe 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Set JAX_PLATFORMS=cpu in GitHub CI
8c159cc 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Tweak build.yml again
df3b907 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Fix build.yml
1e4c8db 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice
Set rounding precision for regression tests
6c16f03 
Signed-off-by: Fabrice Normandin <[email protected]>
@lebrice lebrice merged commit 8c3d69b into master Jun 14, 2024
3 of 4 checks passed
@lebrice lebrice deleted the jax_example branch June 17, 2024 13:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@lebrice