AlphaBetaGamma96

Equinox equivalent of FuncTorch's make_functional

4

Hi @patrick-kidger, Apologises for opening an issue as this is probably something I've missed, but is it at all possible to separate an `eqx.Module` object into a pure function and...

No Batching rules for aten::_linalg_solve_ex, aten::linalg_solve, aten::linalg_solve_ex, and aten::_linalg_slogdet causes significant slowdown for per-sample gradients with torch.linalg.slogdet

16

TL;DR - `torch.linalg.slogdet` is over one order of magnitude slower in computing per-sample gradients in the latest nightly version of PyTorch/FuncTorch (`1.13.0.dev20220721` / ` 0.3.0a0+e8a68f4`) than a previous version of...

Can't vmap(jacrev(slogdet))

17

Hi All, I've been trying to get per-sample gradients of a simple feed-forward network that contains a `torch.slogdet` as its final layer. When I go to apply `vmap` to `jacrev`...

backward hooks aren't called when using vmap

16

Hi All, I've been trying to get the forward activations and backward sensitivies for layers within my model. The `forward_pre_hook` works as expected but `backward_full_hook` fails to get called. A...

memory_efficient_fusion leads to RuntimeError for higher-order gradients calculation. RuntimeError: You are attempting to call Tensor.requires_grad_()

3

Hi All, I've tried improving the speed of my code via using `memory_efficient_fusion`, however, it leads to `Tensor.requires_grad_()` error and I have no idea why. The error is as follows,...

Hi All, I've been trying to use `memory_efficient_fusion` to see if I can speed up a main bottleneck in my code, but I hit a `RuntimeError`. This issue continues from...