[thrtythreeforty]

I don't know JAX, but can this trick be applied as a higher-level function to autodiff-capable languages like JAX?