add jax support

[danielpatrickhug]

https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial7/GNN_overview.html

[danielpatrickhug]

Curious to see if I can dynamically program this with a convert to jax system prompt