Model build uses incorrect input shape, TF 2.16.1 and 2.17 with keras3

[ml-inr]

Suppose that you have a model that splits its input as follows:

x = intuts[:,:,:-1]
mask = intuts[:,:,-1:]

Let further x be used in a custom layer, y=layer(x). This layer should be build using the shape of x. However, since TF version 2.16.1, the layer fails to be build correctly (using keras3, not tf-keras). See this very simple example on google collab.

It was verified that using tf-keras resolves the issue. So that the bug is in keras3 or its compatibility with TF. See also the corresponding issue at TF.

[mehtamansi29]

Hi @ml-inr-

There is conflict for defining name(matrix and bias) in add_weight method. In keras3, add_weight method can accept name like this name=<name of layer>.

By changing the naming declaration in add_weight method will resolve the error.

self.matrix_layer = self.add_weight(name="matrix", shape=(num_fs, self.out_dim), initializer="glorot_uniform")   
self.bias = self.add_weight(name="bias", shape=(self.out_dim,), initializer="zeros")                       

Attached gist for your reference.

[ml-inr]

Hi @mehtamansi29 !
Using keyword name in the add_weight does resolve the issue, thanks! This follows Keras documentation on Layer class, my bad.
Shouldn't, however, my initial code yield an error that I'm not properly using the Layer class (non-keyword arguments)?

[mehtamansi29]

Shouldn't, however, my initial code yield an error that I'm not properly using the Layer class (non-keyword arguments)?

Yes. Your initial code yielding error because declaring name without using keyword. So while building the layer it tries to take "matrix" in self.matrix_layer = self.add_weight("matrix", shape=(num_fs, self.out_dim), initializer="glorot_uniform") as shape.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No