Slightly friendlier import warning.

[koaning]

I initially tried the new keras core with keras_nlp via this code:

import os
os.environ["KERAS_BACKEND"] = "jax"

import keras_nlp

It gave a long trace telling me that "jax" was not installed. Fair enough, so I ran pip install jax and restarted the env. After re-running, it complained that I had to install jaxlib instead.

This is a nit in the grand scheme of things, but it would be nicer if the trace was caught and gave me the correct pip install command. Alternatively, it might also be nice if it told me to python -m pip install keras[jax] so that Keras can handle the jax install on my behalf.

Again, this is a nit, but might improve the developer experience when folks want to give the new Keras Core stuff a spin.

[fchollet]

Thanks for the report -- this is a good idea.

[fchollet]

One complication is that some packages requite a different installation process based on whether you're using GPU or not, and based on your CUDA version.

[fchollet]

We could pick the CUDA version that is compatible with at least one version of each backend and then provide a shell script that installs the right thing (in a way that is mutually compatible) for a given backend, e.g. jax-gpu. We can then point users to this script / command when they hit an ImportError message.