Add softmax layers and convert MNIST example

[hweom]

What does this PR accomplish?

• 🩹 Bug Fix
• 🦚 Feature
• 🧭 Architecture

Changes proposed by this PR:

• Implement Softmax and LogSoftmax layers in the new architecture.
• Convert MNIST example to use the new architecture.
• Fix a bug in native backend where Softmax and LogSoftmax were computed without taking the batch structure into account.

Notes to reviewer:

MNIST example results

Old arch

$ time target/debug/example-mnist-classification mnist linear
Last sample: Prediction: 8, Target: 8 | Accuracy 915/1000 = 91.50%
target/debug/example-mnist-classification mnist linear  3.23s user 0.59s system 97% cpu 3.921 total

$ time target/debug/example-mnist-classification mnist conv
Last sample: Prediction: 8, Target: 8 | Accuracy 946/1000 = 94.60%
target/debug/example-mnist-classification mnist conv  6.48s user 0.93s system 65% cpu 11.307 total

New arch

$ time target/debug/example-mnist-classification mnist linear
Last sample: Prediction: 8, Target: 8 | Accuracy 893/1000 = 89.30%
target/debug/example-mnist-classification mnist linear  2.99s user 0.54s system 97% cpu 3.597 total

$ time target/debug/example-mnist-classification mnist conv
Last sample: Prediction: 8, Target: 8 | Accuracy 943/1000 = 94.30%
target/debug/example-mnist-classification mnist conv  5.80s user 1.30s system 68% cpu 10.398 total

(slight accuracy drop for linear net is likely due to initial weight randomization)

📜 Checklist

• Test coverage is excellent
• All unit tests pass
• The juice-examples run just fine
• Documentation is thorough, extensive and explicit

[drahnr]

A few nits, generally good to go. Variable names are a bit short, especially in iterator closures.

[hweom]

Variable names are a bit short, especially in iterator closures.

Are there specific instances you think are worth improving?

I guess the native backend softmax* and log_softmax* functions? I tried to be minimally invasive there, but maybe it's worth a more thorough cleanup?

[drahnr]

Variable names are a bit short, especially in iterator closures.

Are there specific instances you think are worth improving?

I guess the native backend softmax* and log_softmax* functions? I tried to be minimally invasive there, but maybe it's worth a more thorough cleanup?

It doesn't make sense to do renames here, after a second look I think they're sufficiently contained (usually only spanning up to a handful of lines) that we can leave them as they are to keep the changeset small.