tf.keras.config.set_backend

Reload the backend (and the Keras package).

tf.keras.config.set_backend(
    backend
)

Example:

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

import keras
from keras import ops
type(ops.ones(()))
<class 'tensorflow.python.framework.ops.EagerTensor'>

keras.config.set_backend("jax")
UserWarning: Using `keras.config.set_backend` is dangerous...
del keras, ops

import keras
from keras import ops
type(ops.ones(()))
<class 'jaxlib.xla_extension.ArrayImpl'>

⚠️ WARNING ⚠️: Using this function is dangerous and should be done carefully. Changing the backend will NOT convert the type of any already-instantiated objects. Thus, any layers / tensors / etc. already created will no longer be usable without errors. It is strongly recommended not to keep around any Keras-originated objects instances created before calling set_backend().

This includes any function or class instance that uses any Keras functionality. All such code needs to be re-executed after calling set_backend() and re-importing all imported keras modules.