3 ms·

all issues raised are addressed by jax.vmap

• patrickkidger 1y ago
  Went scrolling looking for this! Most of the article is about problems solved in JAX. Also worth noting the Array API standard exists now. This is generally also trying to straighten out the sharp edges.
  • amensch 1y ago
    Same here, beautiful solution