throws out error when I define a trainable weights as layer

[xiexiaozheng]

I'm trying to use Aimet_tensorflow to quantize a network similar to DETR. One of the steps involves creating a trainable queries vector. While I can directly create a tf.Variable to achieve this in subclass mode, but it becomes a constant parameter when created in the functional model. So, I want to create the following layer class to bypass this limitation.

class QueryEncoding(Layer):
    def build(self, input_shape):
        self.query_embedding = self.add_weight(shape=(16,256), initializer='zeros', dtype=tf.float32, name='query_embedding')
        self.built = True

    def call(self, inputs, training=False):
        return inputs + self.query_embedding

When I use this layer instead of tf.Variable, the prepare_model function throws an error. The error log is from model_preparer.py, which is

The number of weights in the temporary model for unwrapping layer 'query_encoding' does not match the
number of weights of the original layer. The missing weight(s) are {'query_encoding/query_embedding:0'}. This occurs when the Keras Symbolic tensor
passed into the layers call function does not interact with a layer defined inside of the nested layer. Please refer to
the documentation for more information.

This is the call function that is causing this error:
    def call(self, inputs, training=False):
        return inputs + self.query_embedding

do you know how to solve this?

[quic-mangal]

@quic-hitameht, could you help answer this?

[quic-ernst]

@xiexiaozheng Thanks for raising this issue. As I mentioned in #2375, I think this should be updated in the model_preparer. Typically, this error occurs when an internal Keras Symbolic tensor doesn't touch everything in the call function. However, in this case I believe it's because we need to not expand this.

[xiexiaozheng]

@xiexiaozheng Thanks for raising this issue. As I mentioned in #2375, I think this should be updated in the model_preparer. Typically, this error occurs when an internal Keras Symbolic tensor doesn't touch everything in the call function. However, in this case I believe it's because we need to not expand this. @quic-ernst Thank you very much for your response.