meom-group/jaxparrow

`jax.vmap`

Closed this issue · 0 comments

vadmbertr commented

The cyclogeostrophy function should be vmapable.
Currently, some numpy artifacts prevent it:

  • np.ma.masked_array should not be allowed
  • np.* should be replaced by jnp.*