FedProx implementation

[JulienVig]

Closes #802

[JulienVig]

@tomasoignons I rebased the branch onto the latest develop commit. Here is a breakdown of the current commits:

1. tfjs allows training via multiple functions: fit, fitDataset, trainOnBatch. We were using fitDataset on each batch so I switched to using trainOnBatch to be consistent and simplify things
2. We used to compute logits twice for each update, once for the loss and once for the accuracy. This is very inefficient but not super practical to fix. I implemented one but feel free to suggest a better refactoring
3. A temporary commit, I override the training loss by implementing it by hand to then implement the fedprox loss. This commit saves the codebase state where the custom training loss behaves the same as the default loss (make sure this is the case)
4. Fedprox loss building on the previous commit, I didn't have the time to test it in depth so this is up to you

As we were discussing, the fedprox sketch has been hard-coded and doesn't allow users to choose which algorithm they want between fedavg and fedprox, and potential future algorithms like Mime.