MNIST tutorial broken for Colab TPU

[rcrowe-google]

https://colab.sandbox.google.com/github/google/flax/blob/main/docs/nnx/mnist_tutorial.ipynb#scrollTo=6

The MNIST tutorial on the NNX website throws the following error when trying to instantiate the model:

---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
    [... skipping hidden 1 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in _init_backend(platform)
    972   logger.debug("Initializing backend '%s'", platform)
--> 973   backend = registration.factory()
    974   # TODO(skye): consider raising more descriptive errors directly from backend

20 frames
XlaRuntimeError: INVALID_ARGUMENT: Mismatched PJRT plugin PJRT API version (0.47) and framework PJRT API version 0.54).

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
    [... skipping hidden 16 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in backends()
    901           else:
    902             err_msg += " (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)"
--> 903           raise RuntimeError(err_msg)
    904 
    905     assert _default_backend is not None

RuntimeError: Unable to initialize backend 'tpu': INVALID_ARGUMENT: Mismatched PJRT plugin PJRT API version (0.47) and framework PJRT API version 0.54). (set JAX_PLATFORMS='' to automatically choose an available backend)