Richard Zou

functorch doesn't see modules, it only sees PyTorch operations (e.g. add mul ...) so there isn't a way to do this right now

So, it turns out, this isn't sufficient: to do the code generation, we also need to package pytorch/aten/src/ATen/native/native_functions.yaml somehow so that torchgen can access it without requiring a source checkout...

> > aot_function(vmap(f)) does (mostly work), but not the other way around. > > I think this could work, I believe jax recommendation is to jit the vmap and not...

Thanks for the report, we'll take a look soon

Bisected to https://github.com/pytorch/pytorch/pull/75195/files. https://github.com/pytorch/pytorch/pull/75195/files by itself may not be a problem, perhaps the problem is our batching rule for mv. @yueyericardo is the repro you provided the entire model, or...

cc @lezcano @ezyang for https://github.com/pytorch/pytorch/pull/75195 -- this led to a performance regression in functorch. I'm not sure what the original intent of the PR is (there are no tests). I'm...

@yueyericardo - for the original model itself, is the performance regression also 25%, or is it a smaller number? Is the original model public? One thing we can do to...

@Lezcano @ezyang let's say that we did revert the PR (because we're trying to release PyTorch 1.12.1 as soon as possible). Would it cause any other problems? Because the motivation...

> In any case, I'm fine with reverting, but we should investigate what's causing this regardless I agree this warrants more investigation. We've got a problem in that there is...

I patched in https://github.com/pytorch/pytorch/pull/76828 and the above script ends up OOM-ing :/