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>
diff --git a/tools/python/update_helper_files/Dockerfile b/tools/python/update_helper_files/Dockerfile
index 3a4d42c..551674e 100644
--- a/tools/python/update_helper_files/Dockerfile
+++ b/tools/python/update_helper_files/Dockerfile
@@ -8,3 +8,4 @@
 RUN DEBIAN_FRONTEND=noninteractive apt install --no-install-recommends -y python3
 RUN DEBIAN_FRONTEND=noninteractive apt install --no-install-recommends -y pkgconf
 RUN DEBIAN_FRONTEND=noninteractive apt install --no-install-recommends -y libcairo2-dev
+RUN DEBIAN_FRONTEND=noninteractive apt install --no-install-recommends -y libgirepository1.0-dev