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

Question

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

Answer 1

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!

Answer 2

I think it may be that when you use input_shape in the build method of your layer, it means it is a dynamic layer