Add support for JAX

This opens up more exciting solvers, and better optimization than casadi

Change-Id: I824d9ec51c562f0a06b8d2ec1fc27a0e1c8a2ad9
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/tools/python/requirements.txt b/tools/python/requirements.txt
index db0ec26..09121c8 100644
--- a/tools/python/requirements.txt
+++ b/tools/python/requirements.txt
@@ -26,3 +26,5 @@
 tabulate
 
 casadi>=3.6.6
+
+jax[cuda12]