PJit example

### Description of the model to be implemented
ResNet18

### Dataset the model could be trained on
tensorflow_datasets mock_data, e.g.
```
with tfds.testing.mock_data(num_examples=128):
ds = tfds.load('imagenette', split='train')
```

### Specific points to consider

- The existing ImageNet example uses pmap. It is not trivial to replace pmap with pjit in the existing example. A pure pjit example would be great to learn from.
- The example should be functional also on a GPU host for the benefit of the broader community.

### Reference implementations in other frameworks
https://github.com/google/flax/blob/main/examples/imagenet/imagenet.ipynb