Richard Zou

The new functorch 0.2.1 release should fix this problem

Hi @mohamad-amin thanks for the issue and we are happy to help. Can you clarify what exactly is the quantity that you're looking to compute? From reading your messages, I...

> Actually, even your example gives me the same error: (params size is 2 here) > > > RuntimeError: new_args should be of size 1, was of size 2 What...

There were a lot of fixes to functorch w.r.t using `make_functional` and jacobian computation between 1.10 and the nightly. So yes, you'll have to try out the nightly in order...

@mohamad-amin I understand a bit more about neural tangent kernels after reading through some papers. I can probably put together an example in the next week, but I'm curious about...

Hi @mohamad-amin! Thanks for the references to the papers and for the detailed explanation. I'm still working my way through the papers, but I think I have a working example...

Hey @mohamad-amin, I just wanted to give a quick status update (sorry for leaving you hanging!). We've been stuck on a bug https://github.com/pytorch/functorch/issues/417 that makes the NTK computation error out...

@mohamad-amin we finally fixed https://github.com/pytorch/functorch/issues/417! Here's the first version of the example of how to compute NTKs with functorch. https://github.com/pytorch/functorch/blob/cb876fad2b2a9269424c8212a82652f372db6dfa/notebooks/neural_tangent_kernels.ipynb . The code runs on the latest build of functorch...

Originally reported at https://discuss.pytorch.org/t/torch-vmap-batching-rule-not-implemented-for-aten-item-we-could-not-generate-a-fallback/144407

Okay, I just realized what is happening here. The problem isn't advanced indexing, the problem is we are vmapping over a Tensor and using it to index into a Python...