tree 9faff2b9a82d0623f976be9d683709e9b8ef3d23
parent 9e6951254f935c2a6379c2140a071a8ae7552f20
author Austin Schuh <austin.linux@gmail.com> 1725314556 -0700
committer Austin Schuh <austin.linux@gmail.com> 1726961689 -0700
gpgsig -----BEGIN PGP SIGNATURE-----
 
 iQGzBAABCgAdFiEExcNA7Sy+j6Ly5hA4QVYw08sOnWYFAmbvWBoACgkQQVYw08sO
 nWa7Pgv/dJvuzP4dk4wKQdCPo/tSKuP1r+OKAc6Y8B6h8L8F221BM8kzmJontFtB
 MJoxy6xNgl8DOdxBs0E/wv553E7s1TyG5Gh9o3LPc14/yjPRoKYhUpyh5sX2rnQQ
 kepMilC2vVkZf7NA0/Yl4MmCr0onn2X/8hFA/1n7C5FCjYzJzjqyOQzeYm5CmcHn
 8DwW9RjM/4/F8ms/QZiGt7l4nHmxjaGXAlCbVvJFY7zMb86WlQ51fu5nlo1/gVHs
 Po5eQcm95cg1usU+OlTnCWksovWpPO1m4rHdless+QFxcODgYZ+06VAcm+T1iz93
 srnS8Z5pUcTdvSxQLBz64FVth96xGVD2sSrxGNVm4FdrwJ2AzjP01jFNQ2mZnyVS
 2zQp77raa/ZHZMp6TvetOjVXTOaqIA30y6cahxaLn/eiagtG3X06QQlxkonzJyVs
 uVcDbvOYSyP/FBS3Qw76z7OHcbuDqppafSuFxUgScU2TNfS6GQIDT88l/kkVrr7C
 36wnszaU
 =jEAz
 -----END PGP SIGNATURE-----

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>