search

talmolab / stac-mjx

Implementation of STAC using MJX for GPU acceleration. Part of VNL project.
BSD 3-Clause "New" or "Revised" License
4 stars 2 forks source link

Change jaxopt usage to optax #28

Open charles-zhng opened 1 month ago

charles-zhng commented 1 month ago

jaxopt is being merged with optax. All the features we use are in optax anyway, so reimplement using the optax-native functions instead.

charles-zhng commented 3 weeks ago

Optax projections here: https://optax.readthedocs.io/en/latest/api/projections.html