Skip to content

Avoid using vmap when parallel_chunk_size=1#221

Closed
ValerianRey wants to merge 8 commits intomainfrom
add-rnn-test
Closed

Avoid using vmap when parallel_chunk_size=1#221
ValerianRey wants to merge 8 commits intomainfrom
add-rnn-test

Conversation

@ValerianRey
Copy link
Copy Markdown
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 added sandbox cc: test Conventional commit type for changes to tests. package: autojac labels Dec 27, 2024
@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
@codecov
Copy link
Copy Markdown

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 cc: docs Conventional commit type for changes to the documentation. cc: feat Conventional commit type for new features. 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

cc: docs Conventional commit type for changes to the documentation. cc: feat Conventional commit type for new features. cc: test Conventional commit type for changes to tests. package: autojac

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant