Skip to content

VMAP JVP in Hutchinson#23

Open
stevet40 wants to merge 1 commit into
mainfrom
trace
Open

VMAP JVP in Hutchinson#23
stevet40 wants to merge 1 commit into
mainfrom
trace

Conversation

@stevet40

@stevet40 stevet40 commented Jun 2, 2026

Copy link
Copy Markdown
Member

This copies the VMAP JVP trick from the Hutch++ and XTrace into the Hutchinson part of the forward. Also implements a flexible number of probe vectors for Hutchinson (fixes #21 ). There is probably some redundancy / cleanup possible with this, which should maybe be a separate issue.

@gurjeetjagwani

Copy link
Copy Markdown
Member

This looks good, but we ought to benchmark against normal autograd especially for our analysis reported in StAD. Because I think this must be faster?

@stevet40

stevet40 commented Jun 8, 2026

Copy link
Copy Markdown
Member Author

Yes, I think we should definitely revisit the implementation of this before we tackle #13, to see if there is any room for improvement.

@gurjeetjagwani

Copy link
Copy Markdown
Member

I will run an analysis and profiling of this code against the old and share the results soon. Once, we have the results we can identify places where we can improve the implementation and finalise the code.

@gurjeetjagwani

Copy link
Copy Markdown
Member

Does that sound reasonable?

@stevet40

stevet40 commented Jun 8, 2026

Copy link
Copy Markdown
Member Author

Sounds good! Thanks!!

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Option to specify number of probe vectors for vanilla Hutchinson

2 participants