MNIST training example in FP8

[balancap]

We need to validate AutoScale FP8 training on the basic MNIST example to start with.

With multiple experiments:

• Fwd FP8-143 format used;
• Bwd FP8-152 format used;
• Fwd + bwd;

These experiments can be done relying primarily on JAX LAX reduce_precision operator. No need of FP8 hardware for now.

[balancap]

cc @lyprince in case you're interested :)

[balancap]

Basic MNIST FP8 training implemented in #87
Training is working as expected, even though hard to make any conclusion from just this simple MNIST model.