Support other fitting methods #1
Description
It might be useful to support other fitting methods, especially SGD e.g. based on https://optax.readthedocs.io/en/latest/. This might help to improve generalization and support for mini-batches when fitting.