Add Pytorch or Jax examples

[jwuphysics]

Thanks for launching this site! I love the examples across different astronomical domains using different kinds of data-driven models!

I figured that it could be nice to include more examples using Pytorch and Jax. Both of these frameworks are seeing strong adoption in research domains, and it's likely that Jax will supersede TensorFlow at some point (if it has not already).

[cmurray-astro]

This sounds like an excellent idea, thanks John! We will include it on our list for the future, and are looking forward to incorporating your recently-submitted example!