torchjd

torchjd copied to clipboard

TODO: - [ ] Add a special type of error for that, like UnsupportedOnSparseError. We could catch that + NotImplementedError, but not catch ValueError for instance. - [ ] Prioritize...

ValerianRey

feat(autogram): Use new engine

1

Already making the PR to keep track of progress, and to be able to update the dev-new-engine branch easily when main is updated. - [x] Remove explicit batched optimizations #470...

ValerianRey

Add an automatic tagger to the CI

2

Note that this does not remove tags if not needed anymore. I think this shouldn't in case we want to manually add a tag, for instance if we change documentation...

PierreQuinton

refactor(autogram): Remove vmap rule generation

3

Not sure what I'm doing here, but as far as I understand, the generated vmap rule for JacobianAccumulator is never used. I think (correct me if I'm wrong) that it...

ValerianRey

feat(autogram): Do not consider cross-terms

2

* Add and use JacobianBasedGramianComputerWithoutCrossTerms * Make ModelAlsoUsingSubmoduleParamsDirectly and InterModuleParamReuse xpass * Fix missing type hint * Test against no cross-terms

ValerianRey

feat(autogram): Add `ModuleBasedGramianComputer`.

5

PierreQuinton

I think we need some script in `.github/workflow` and a branch protection. It would make the `TODO` keyword very useful and safe.

PierreQuinton

Here are some thoughts on how we could support randomness in `autogram`. Consider this context: ```python class local_rng_context: """ A context manager that saves the global CPU and CUDA RNG...

PierreQuinton

Sparse MatMul Support in TorchJD for Seamless compatibility with torch.vmap()

5

Hi folks, it's me again, I had noticed a few months back that TorchJD bugged when trying to input sparse matrices into the custom JD backward pass due to incompatibility...

EmileAydar

Will this repo support `jax` framework?

2

Thanks for your solid work! Since `jax` has been widely used, will this repo offer a `jax` version?

HiccupRL