Add a lot of optimizers including optax support to sgd
#2041
+300
−43
optax support to sgd
#2041