Cannot convert model with dynamic variable

[ismaeelbashir03]

I am writing a matrix multiplication accumulator that sums the results of calls to the model to a variable:

class AccumulateMatrixMultiplication(tf.Module):
        def __init__(self):
            super().__init__()
            self.accumulator = None

        @tf.function(input_signature=[
            [
                tf.TensorSpec(shape=[1, None, None], dtype=tf.float32),
                tf.TensorSpec(shape=[1, None, None], dtype=tf.float32),
                tf.TensorSpec(shape=[1], dtype=tf.bool)
            ]
        ])
        def call(self, inputs):
            matrix_a, matrix_b, reset_flag = inputs
            
            # Initialize or reset a dynamic accumulator
            if self.accumulator is None or tf.reduce_any(reset_flag):
                self.accumulator = tf.Variable(
                    initial_value=tf.zeros(matrix_a.shape, dtype=tf.float32),
                    trainable=False
                )

            # Compute result
            result = tf.matmul(matrix_a, matrix_b)

            self.accumulator = self.accumulator.assign_add(result[tf.newaxis, ...])

            return self.accumulator

However, this gives:

ValueError: in user code:

        self.accumulator = tf.Variable(

    ValueError: Cannot convert a partially known TensorShape (1, None, None) to a Tensor.

[gaikwadrahul8]

Hi, @ismaeelbashir03
Thank you for reporting this issue. I have been able to reproduce similar behavior after making minor adjustments to the provided code snippet. However, I am also encountering the same error you described. Please refer this gist-file so we'll have to dig more into this issue and will update you.

Please find the error log below :

ValueError: in user code:

    File "<ipython-input-12-8cb3748a0e9c>", line 15, in __call__  *
        self.accumulator = tf.Variable(

    ValueError: Cannot convert a partially known TensorShape (1, None, None) to a Tensor.

Thank you for your cooperation and understanding.