commit ea5f0a7624d44a82c7425cdf4909f2bd722d7f49
author Austin Schuh <austin.linux@gmail.com> Mon Sep 02 15:02:36 2024 -0700
committer Austin Schuh <austin.linux@gmail.com> Sat Sep 21 16:34:49 2024 -0700
tree 9faff2b9a82d0623f976be9d683709e9b8ef3d23
parent 9e6951254f935c2a6379c2140a071a8ae7552f20

Add flax and friends for training neural networks

This updates our formatter, which changed what it expects very slightly.
Update all the pieces of code to match

optax is for the actual optimizers, jax is the underlying accelerated
linear algebra library, tensorflow is for loading datasets and exporting
models.

Change-Id: Ic4c3b425cda74267e1d0ad1615c42452cbefab8a
Signed-off-by: Austin Schuh <austin.linux@gmail.com>

tools/python/requirements.txt

diff --git a/tools/python/requirements.txt b/tools/python/requirements.txt
index 09121c8..5d2234d 100644
--- a/tools/python/requirements.txt
+++ b/tools/python/requirements.txt
@@ -8,7 +8,7 @@
 opencv-python
 osqp
 pkginfo
-pygobject
+pygobject==3.42.2
 requests
 scipy
 shapely
@@ -28,3 +28,12 @@
 casadi>=3.6.6
 
 jax[cuda12]
+optax
+flax
+clu
+
+# Pin this so we can actually depend on it.  Used by tensorflow/jax.
+setuptools==74.0.0
+
+tensorflow
+tensorflow_datasets