[FIXED] Is it still necessary to implement `compute_output_shape()` when defining a custom tf.keras Layer?

Issue

I have implemented a custom Layer in tf.keras, using TensorFlow 2.1.0.

In the past, when using the stand-alone Keras, it was important to define the compute_output_shape(input_shape) method in any custom layer so that the computational graph could be created.

Now, having moved to TF2, I found out that even if I remove that method from my custom implementation the layer still works as expected. Apparently, it works both in eager and graph mode. This is an example of what I mean:

from tensorflow.keras.layers import Layer, Input
from tensorflow.keras.models import Sequential
import numpy as np


class MyLayer(Layer):
    def call(self, inputs):
        return inputs[:, :-1]  # Do something that changes the shape


m = Sequential([MyLayer(), MyLayer()])
m.predict(np.ones((10, 3)))  # This would not have worked in the past

Is it safe to say that compute_output_shape() is not necessary anymore? Am I missing something important?

In the documentation there's no explicit mention of removing compute_output_shape(), although none of the examples implements it explicitly.

Thanks

Solution

It is not mentioned in the Tensorflow Documentation but in Chapter 12, Custom Models and Training with TensorFlow of the book, Hands-on Machine Learning using Scikit-Learn and Tensorflow (2nd Edition Updated for Tensorflow 2) of O'REILLY Publications, written by Aurelien Geron, it is mentioned as shown in the screenshot below:

To answer your question, yes, it is safe to say compute_output_shape is not needed unless the Layer is Dynamic.

This is evident from this Tensorflow Tutorial on Custom Layer where compute_output_shape is not used.

Hope this helps. Happy Learning!

Answered By - Tensorflow Warrior