Leveraging Python and JAX in R workflows