TensorFlow.js: Train MNIST with the Layers API