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

Avoid using vmap when parallel_chunk_size=1 #221

Closed
wants to merge 8 commits into from
Closed

Conversation

ValerianRey
Copy link
Contributor

@ValerianRey ValerianRey commented Dec 27, 2024

  • Add a special case to Jac when parallel_chunk_size=1 + comment to explain that this is a fix to vmap
  • Add tests of Jac with chunk_size=1.
  • Add tests of backward and mtl_backward with parallel_chunk_size=1.
  • Add test with cyclic tensor graph
  • Add changelog entry

Note that even without the change to Jac, and even with CUDA, test_tensor_used_multiple_times succeeds, so the issue is really with CUDA + RNN rather than with CUDA + the graph of tensors being cyclical. Still, I think this test is interesting so I kept it.

@ValerianRey ValerianRey self-assigned this Dec 27, 2024
@ValerianRey ValerianRey changed the title Add test_rnn in test_backward.py Add test for RNN Dec 27, 2024
Copy link

codecov bot commented Dec 27, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Files with missing lines Coverage Δ
src/torchjd/autojac/_transform/jac.py 91.83% <100.00%> (+1.36%) ⬆️

@ValerianRey ValerianRey changed the title Add test for RNN Improve RNN support Dec 29, 2024
@ValerianRey ValerianRey added docs Improvements or additions to documentation feat New feature or request labels Dec 29, 2024
@ValerianRey ValerianRey force-pushed the add-rnn-test branch 2 times, most recently from 4c279aa to d699f6d Compare December 29, 2024 21:09
@ValerianRey ValerianRey changed the title Improve RNN support Avoid using vmap when parallel_chunk_size=1 Dec 29, 2024
@ValerianRey ValerianRey deleted the add-rnn-test branch December 29, 2024 22:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
docs Improvements or additions to documentation feat New feature or request package: autojac sandbox test
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant