Minimal MNIST workload

[znado]

A rough sketch of an implementation of our current API as outlined in the rules draft, using jax/numpy.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[mikerabbat]

This is a great start! I think it covers almost all of the main points. I've added minor comments inline and I have three high-level comments.

The main thing that seems to be missing is a way for submissions to specify the hyperparameter search space. Would it make sense for submissions to provide one more function that accepts a source of randomness and outputs a sample from the tuning search space?

The way PRNG seeds are handled here is pretty specific to Jax. Different frameworks will probably need to do this differently. Is there any utility to making the harness code as framework-agnostic as possible? Is there any concern that framework-specific differences in how PNRGs are handled can become a non-trivial source of difference between implementations?

We've been discussing the benchmark in terms of a system that has 8 GPUs. Will submissions have control over how those GPUs are used? For some workloads (e.g., a ResNet50/ImageNet-like one) it may be useful to do data-parallel training using all 8 GPUs. For others (smaller ones like CIFAR-10), that may be overkill and a submission may want to try something different. The current API doesn't seem to expose this flexibility.