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

[functorch] jacrev: chunk_size=1 without vmap #91326

Closed
wants to merge 6 commits into from
Closed

[functorch] jacrev: chunk_size=1 without vmap #91326

wants to merge 6 commits into from

Conversation

kshitij12345
Copy link
Collaborator

@kshitij12345 kshitij12345 commented Dec 22, 2022 •
edited by pytorch-bot bot

@pytorch-bot
Copy link

pytorch-bot bot commented Dec 22, 2022 •
edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/91326

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Failures

As of commit f6d92f7:

BROKEN TRUNK - The following jobs failed but were present on the merge base 06bdd49:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@kshitij12345 kshitij12345 marked this pull request as ready for review December 22, 2022 21:08
zou3519
zou3519 reviewed Dec 22, 2022
View reviewed changes
Copy link
Contributor

@zou3519 zou3519 left a comment •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor comment on edge case and testing and docs, but aside from that, this LGTM

torch/_functorch/eager_transforms.py Show resolved Hide resolved
test/functorch/test_eager_transforms.py Outdated Show resolved Hide resolved
torch/_functorch/eager_transforms.py
@@ -492,9 +495,23 @@ def compute_jacobian_stacked():
for flat_basis_chunk in _chunked_standard_basis_for_(flat_output,
flat_output_numels,
chunk_size=chunk_size):
if chunk_size == 1:
flat_basis_chunk = tree_map(lambda t: torch.squeeze(t, 0), flat_basis_chunk)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should this be tree_map_only? (what happens for non-Tensor arguments?)

Also, squeeze is a bit dangerous, because if you squeeze on a dimension of size 2 it doesn't actually do anything. Would prefer an assertion that t.shape(0) == 1

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These flat_basis_chunk are guaranteed to be tensors. Will add an assert.

Ref:

pytorch/torch/_functorch/eager_transforms.py

Line 615 in f62a3ca

def _chunked_standard_basis_for_(tensors, tensor_numels, chunk_size=None):

test/functorch/test_eager_transforms.py Outdated Show resolved Hide resolved
@drisspg drisspg added triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module module: functorch Pertaining to torch.func or pytorch/functorch labels Dec 23, 2022
@kshitij12345 kshitij12345 added the release notes: functorch release notes category; Pertaining to torch.func or pytorch/functorch label Dec 23, 2022
@kshitij12345
Copy link
Collaborator Author

Have addressed the reviews. PTAL @zou3519 :)

zou3519
zou3519 approved these changes Dec 27, 2022
View reviewed changes
Copy link
Contributor

@zou3519 zou3519 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thank you!

@kshitij12345
Copy link
Collaborator Author

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Dec 27, 2022
@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: This PR is too stale; the last push date was more than 3 days ago. Please rebase and try again. You can rebase by leaving the following comment on this PR:
@pytorchbot rebase

Details for Dev Infra team Raised by workflow job

@kshitij12345
Copy link
Collaborator Author

@pytorchbot merge -f"CI failure is unrelated"

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@zou3519 zou3519 zou3519 approved these changes

Assignees
No one assigned
Labels
ciflow/trunk Trigger trunk jobs on your pull request Merged module: functorch Pertaining to torch.func or pytorch/functorch open source release notes: functorch release notes category; Pertaining to torch.func or pytorch/functorch triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
5 participants
@kshitij12345 @pytorchmergebot @zou3519 @pytorchbot @drisspg