flax
Makes annotated_mnist use Optax's xent loss.
#2071
Merged

Loading