MNIST tutorial broken for Colab TPU

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)
```