Tensorflow compute_output_shape() Not Working For Custom Layer - tensorflow

Question

I have created a custom layer (called GraphGather) in Keras, yet the output tensor prints as :
Tensor("graph_gather/Tanh:0", shape=(?, ?), dtype=float32)
For some reason the shape is being returned as (?,?), which is causing the next dense layer to raise the following error:
ValueError: The last dimension of the inputs to Dense should be defined. Found None.
The GraphGather layer code is as follows:
class GraphGather(tf.keras.layers.Layer):
def __init__(self, batch_size, num_mols_in_batch, activation_fn=None, **kwargs):
self.batch_size = batch_size
self.num_mols_in_batch = num_mols_in_batch
self.activation_fn = activation_fn
super(GraphGather, self).__init__(**kwargs)
def build(self, input_shape):
super(GraphGather, self).build(input_shape)
def call(self, x, **kwargs):
# some operations (most of def call omitted)
out_tensor = result_of_operations() # this line is pseudo code
if self.activation_fn is not None:
out_tensor = self.activation_fn(out_tensor)
out_tensor = out_tensor
return out_tensor
def compute_output_shape(self, input_shape):
return (self.num_mols_in_batch, 2 * input_shape[0][-1])}
I have also tried hardcoding compute_output_shape to be:
python
def compute_output_shape(self, input_shape):
return (64, 150)
```
Yet the output tensor when printed is still
Tensor("graph_gather/Tanh:0", shape=(?, ?), dtype=float32)
which causes the ValueError written above.
System information
Have written custom code
**OS Platform and Distribution*: Linux Ubuntu 16.04
TensorFlow version (use command below): 1.5.0
Python version: 3.5.5

Answer 1

I had the same problem. My workaround was to add the following lines to the call method:
input_shape = tf.shape(x)
and then:
return tf.reshape(out_tensor, self.compute_output_shape(input_shape))
I haven't run into any problems with it yet.

Answer 2

If Johnny's answer doesn't work, I found another way to get around this is to follow advice here https://github.com/tensorflow/tensorflow/issues/38296#issuecomment-623698709
which is to call the set_shape method on the output of your layer.
E.g.
l=GraphGather(...)
y=l(x)
y.set_shape( l.compute_output_shape(x.shape) )
This only works if you are using the functional API.