Why stealthy Tensorflow upgrade to 2.9.1 breaks Google Colab Jupyter notebooks that use GPU but not CPU?

Sometime in the last three days, an upgrade to TensorFlow on Google colab was made from 2.8.x to 2.9.1. This upgrade broke all of my current research notebooks including a minimal MNIST example I have included. A thorough review of the release notes does not show that any of the packages that I am using in Keras or TensorFlow were changed.

Further work on this error revealed it only occurs when the Colab runtime includes a GPU. It works fine on a Colab CPU or TPU. Here is a 34 line example that reproduces the error:

    import tensorflow as tf
    import keras
    
    # the data, split between train and test sets
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    
    x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
    x_test  = x_test.reshape ( x_test.shape[0], 28, 28, 1)
    input_shape = (28, 28, 1)
    
    x_train  = x_train.astype('float32')
    x_test   = x_test.astype('float32')
    x_train /= 255
    x_test  /= 255
    
    # convert class vectors to binary class matrices
    y_train = keras.utils.to_categorical(y_train, 10)
    y_test  = keras.utils.to_categorical(y_test , 10)
    
    model = keras.models.Sequential()
    model.add(keras.layers.Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape))
    model.add(keras.layers.Conv2D(64, (3, 3), activation='relu'))
    model.add(keras.layers.MaxPooling2D(pool_size=(2, 2)))
    model.add(keras.layers.Dropout(0.25))
    model.add(keras.layers.Flatten())
    model.add(keras.layers.Dense(128, activation='relu'))
    model.add(keras.layers.Dropout(0.5))
    model.add(keras.layers.Dense(10, activation='softmax'))
    
    model.compile(loss=keras.losses.categorical_crossentropy, optimizer=keras.optimizers.Adam(),
                  metrics=['accuracy'])
    
    model.fit(x_train, y_train, batch_size=100, epochs=1, verbose=1, validation_data=(x_test, y_test))
    score = model.evaluate(x_test, y_test, verbose=0)
    print('Test loss:', score[0], 'Test accuracy:', score[1])

I have included a trace of the errors below.

When I downgrade back to tensorflow 2.8.2 the error disappears and all my collab notebooks work normally.

To keep using a GPU, the current workaround, which adds 86 seconds per run is:

    !pip install tensorflow==2.8.2
    import tensorflow as tf
    print(tf.__version__)

• Error Log when GPU is included in runtime configuration:

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11490434/11490434 [==============================] - 0s 0us/step
---------------------------------------------------------------------------
UnimplementedError                        Traceback (most recent call last)
[<ipython-input-1-05f207168698>](https://localhost:8080/#) in <module>
     31               metrics=['accuracy'])
     32 
---> 33 model.fit(x_train, y_train, batch_size=100, epochs=1, verbose=1, validation_data=(x_test, y_test))
     34 score = model.evaluate(x_test, y_test, verbose=0)
     35 print('Test loss:', score[0], 'Test accuracy:', score[1])

1 frames
[/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py](https://localhost:8080/#) in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     53     ctx.ensure_initialized()
     54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:
     57     if name is not None:

UnimplementedError: Graph execution error:

    [...]

Node: 'sequential/conv2d/Conv2D'
DNN library is not found.
     [[{{node sequential/conv2d/Conv2D}}]] [Op:__inference_train_function_865]