AlphaBetaGamma96

So it seems the solution is to place the `vmap` called within the `memory_efficient_fusion` call like so, ``` ps_elocal_fusion = memory_efficient_fusion(vmap(grad(local_kinetic_from_log_vmap, argnums=0), in_dims=(None, 0))) ``` then just call, ``` ps_elocal_fusion(params,...

I thought I'd also mention that `memory_efficient_fusion` fails if a scalar is included. For example, using this function (which differs from the original value of `-0.5 *` factor) ``` def...

What version of functorch are you using? Your minimal reproducible script works for me. It returns, ``` tensor([[[1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], [0., 1.,...

Hi @patrick-kidger, this is exactly what I was looking for! Thank you for the clear example code! If `jit-vmap-grad` computes per-sample gradients, I assume `jit-grad-vmap` computes loss-gradients? (or something else?)...

Hi @zou3519, I indeed am! Shall I upgrade to the latest nightly?

Hi @zou3519, I do have an example in mind. I'll have a look for it and share a minimal reproducible example below as soon as I find it!

Hi @zou3519, Sorry for this being quite late, but I had to rewrite it. I've attached an example below, which is run with the current versions of torch = `1.13.0.dev20221003`...

Hi @zou3519, any update on this?

Hi @zou3519, apologies for the late response. The refactor of the example looks great! Can I ask a few questions? 1. Is the feature flag represented as the class attribute...

Apologizes for the late reply @zou3519 ! > We have docs now, that should answer the questions :) https://pytorch.org/docs/master/notes/extending.func.html I shall give this a read! > Yes, the syntax is...