Avoid re-flattening in jit() and pmap() when no donate_argnums are present. #3945
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Following the same special-casing of static_argnums, this should provide a speedup specially when the number of arguments provided is large.
tomhennigan
approved these changes
Aug 3, 2020
There was a problem hiding this comment.
BTW if you were seeing this overhead when using Haiku you should find google-deepmind/dm-haiku@81f6f7a significantly improves the performance for you (tree_flatten on params/state is now O(1)).
|
Thank you! |
|
@tomhennigan This PR fails some internal tests. Perhaps you can take a look and see what the problem is. |
|
The failure is in BufferDonationTest.test_pmap_donate_argnums_invalidates_input (only on TPU and GPU): |
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Aug 4, 2020
…nt. (google#3945)" This reverts commit 4e873f4. See comments in google#3945 about the failure.
gnecula
added a commit
that referenced
this pull request
Aug 4, 2020
|
I see, I hadn't noticed that donate_argnums can be an integer too (which is
the case in that test), so the condition can't be using the implicit
boolean value.
…On Tue, 4 Aug 2020, 12:10 George Necula, ***@***.***> wrote:
The failure is in
BufferDonationTest.test_pmap_donate_argnums_invalidates_input:
File "/build/work/978904e12079446cd55d2263ccdfe9103bb3/google3/runfiles/google3/third_party/py/jax/tests/api_test.py", line 3391, in test_pmap_donate_argnums_invalidates_input
self.assertDeleted(x)
File "/build/work/978904e12079446cd55d2263ccdfe9103bb3/google3/runfiles/google3/third_party/py/jax/tests/api_test.py", line 3404, in <lambda>
assertDeleted = lambda self, x: self._assertDeleted(x, True)
File "/build/work/978904e12079446cd55d2263ccdfe9103bb3/google3/runfiles/google3/third_party/py/jax/tests/api_test.py", line 3412, in _assertDeleted
self.assertEqual(buffer.is_deleted(), deleted)
File "<embedded stdlib>/unittest/case.py", line 829, in assertEqual
assertion_func(first, second, msg=msg)
File "<embedded stdlib>/unittest/case.py", line 822, in _baseAssertEqual
raise self.failureException(msg)
AssertionError: False != True
—
You are receiving this because you authored the thread.
Reply to this email directly, view it on GitHub
<#3945 (comment)>, or
unsubscribe
<https://github.com/notifications/unsubscribe-auth/AAINVQYFJZLCC4LVRNLHLWDR67NAHANCNFSM4PTSJQXA>
.
|
|
PR resent #3955. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Following the same special-casing of static_argnums, this should provide a speedup specially when the number of arguments provided is large.