Skip to content

Switch from flax.optim to optax

Vicentini Filippo requested to merge optax into nk3

Created by: PhilipVinc

This also supports haiku out of the box. See the included example @femtobit

Merge request reports