Suryanarayanan Manoj Sanu

PETSc with original ad_wrapper is computing wrong gradients

1

This has to be fixed! We may need to keep short gradient checks for each function to detect where the leak is happening.

For now, we have used `device_put()` to put the global stiffness matrix into CPU. Need to create a transfer guard to ensure that this is indeed happening.

Support for custom solvers

1

Hi guys, I am using a JAX-based framework to do finite element analysis. Currently, I use solvers from two sources; either `jax.scipy.bicgstab` or PETSc. I like the fact that `lineax`...

- PETSc has direct solvers - We can easily implement them

- To make it easier to run smaller problems, it might be beneficial to implement analytical formulas for simple objectives like compliance

- According to Tianju, these are named wrong

- [x] General cleanup + Flake8 file -> Surya - [x] Enable `pip install` in editable mode [Setup] -> Surya - [x] Environment file -> Surya - [x] All tests...

Multi-step loading + auto-diff

2

- In the plasticity example, the load is applied in smaller steps. This is beneficial for easy convergence of iterative solvers for other cases as well. - Solution: - Just...

From https://github.com/tianjuxue/jax-am/issues/28

- Currently, the solver [when using petsc] cannot be combined with any of Jax's transforms - This prevents vmapping and higher order derivatives - Solution: - Wrap petsc with Lineax...