flexcompute / tidy3d

Fast electromagnetic solver (FDTD) at scale.
https://docs.flexcompute.com/projects/tidy3d/en/latest/
GNU Lesser General Public License v2.1
179 stars 40 forks source link

PyTorch (and maybe JAX) plugins for wrapping objective functions #1945

Open yaugenst-flex opened 3 weeks ago

yaugenst-flex commented 3 weeks ago

We need some plugins to convert objective functions defined using autograd (since that is what we use) to PyTorch and possibly JAX. Proposed structure:

tidy3d/plugins
├── pytorch/
└── jax/

For JAX, we can probably just use agjax