TensorFlow Keras

Keras is an open-source python library which provides many common building blocks to ease development of deep neural network code.


In the past Keras was a separate project. Currently Keras is part of TensorFlow, available as tf.keras module. This is the only Keras version supported on Gaudi.

Keras API Support

The following Keras APIs are supported on Gaudi:

  • tf.keras.activations.*,

  • tf.keras.applications.*,

  • tf.keras.backend.*,

  • tf.keras.callbacks.*,

  • tf.keras.constraints.*,

  • tf.keras.estimator.*,

  • tf.keras.initializers.*,

  • tf.keras.layers.*,

  • tf.keras.losses.*,

  • tf.keras.metrics.*,

  • tf.keras.mixed_precision.*,

  • tf.keras.models.*,

  • tf.keras.optimizers.*,

  • tf.keras.regularizers.*,

  • tf.keras.utils.*,

  • tf.keras.wrappers.*,

The following APIs can be used, but some operations may be delegated to CPU:

  • tf.keras.datasets.*,

  • tf.keras.preprocessing.*,

  • all experimental APIs including tf.keras.experimental.*,


tf.keras.mixed_precision is the recommended mixed precision mechanism for Keras models on Gaudi. To start using tf.keras.mixed_precision, set the mixed_bfloat16 policy and float32 data type for a last layer in a model as described in the TensorFlow Mixed Precision Guide.


tf.keras.applications contains several models that can be used “as is” with pre-trained weights, used as a base or trained from scratch.


Training from scratch was verified only on limited number of models from tf.keras.applications.


tf.keras.optimizers contains various optimizers that can be used in the models. In TensorFlow 2.11, the previously available optimizers were moved to tf.keras.optimizers.legacy namespace. Consequently, tf.keras.optimizers contains now a set of optimizers based on the new tf.keras.optimizers.Optimizer base class that differs in API and functionality.

To ease migration and maintain models compatibility between TensorFlow 2.8 and 2.11, tf.keras.optimizers should be replaced with habana_frameworks.tensorflow.backward_compatible_optimizers class which is also available in Habana® Gaudi® TensorFlow library. It allows to transparently use tf.keras.optimizers namespace for TensorFlow 2.8, and tf.keras.optimizers.legacy namespace for TensorFlow 2.11.

For example, when using SGD optimizer, to leverage the approach allowing for compatibility between TensorFlow versions, replace test_optimizer = tf.keras.optimizers.SGD(0.1) with the following code lines:

from habana_frameworks.tensorflow import backward_compatible_optimizers
test_optimizer = backward_compatible_optimizers.SGD(0.1)


Using habana_frameworks.tensorflow.backward_compatible_optimizers class is not necessary if maintaining the model compatibility between TensorFlow 2.8 and 2.11 is not required. In this case, tf.keras.optimizers can be replaced directly with tf.keras.optimizers.legacy.