search

meom-group / jaxparrow

A package for computing the inversion of the cyclogeostrophic balance based on a variational formulation approach, using JAX.
https://jaxparrow.readthedocs.io/
Apache License 2.0
7 stars 1 forks source link

lax.scan #64

Closed vadmbertr closed 5 months ago

vadmbertr commented 5 months ago

replace lax.while_loop with lax.scan as the latter is rev-mode differentiable