Move Jax import to conditional

[sampathweb]

TensorFlow Nightly failed on import of Keras since jax becomes a requirement due to import.

+ /tmp/tmp.yijSXl7avZ/bin/python3 -c 'import sys; import tensorflow as tf; sys.exit(0 if "keras" in tf.keras.__name__ else 1)'
Traceback (most recent call last):
  File "/tmp/tmp.yijSXl7avZ/lib/python3.12/site-packages/tensorflow/python/util/lazy_loader.py", line 147, in _initialize
    import keras  # pylint: disable=g-import-not-at-top
    ^^^^^^^^^^^^
  File "/tmp/tmp.yijSXl7avZ/lib/python3.12/site-packages/keras/__init__.py", line 8, in <module>
    from keras import _tf_keras
  File "/tmp/tmp.yijSXl7avZ/lib/python3.12/site-packages/keras/_tf_keras/__init__.py", line 1, in <module>
    from keras._tf_keras import keras
  File "/tmp/tmp.yijSXl7avZ/lib/python3.12/site-packages/keras/_tf_keras/keras/__init__.py", line 61, in <module>
    from keras._tf_keras.keras import layers
  File "/tmp/tmp.yijSXl7avZ/lib/python3.12/site-packages/keras/_tf_keras/keras/layers/__init__.py", line 148, in <module>
    from keras.src.utils.jax_layer import FlaxLayer
  File "/tmp/tmp.yijSXl7avZ/lib/python3.12/site-packages/keras/src/utils/jax_layer.py", line 3, in <module>
    import jax
ModuleNotFoundError: No module named 'jax'

JaxLayer checks for backend, so moving this import should be no-op, but removes the need to have it installed in a TF backend.