Does TensorFlow recalculate tensor if inputs do not change?

Viewed 16

Consider this simple example, where we execute the function get_d multiple times in a tf.scan call. get_d calls get_c, which depends on model parameters, which do not change during the tf.scan iteration.

import tensorflow as tf

class Testclass(tf.Module):
    def __init__(self):
        # Some model variables
        self.a = tf.Variable(
            tf.random.uniform([200,3000]),
            trainable=True,
        )
        self.b = tf.Variable(
            tf.random.uniform([200,3000]),
            trainable=True,
        )
    
    # Calculate something from the model parameters
    def get_c(self):
        return tf.transpose(self.a) @ self.b
    
    # Frequently used function
    def get_d(self, x, v):
        c = self.get_c()  # This is required since the self.a and self.b evolve over time
        return x + tf.linalg.matvec(c,v)  # Just some calculation based on c

# Some sample model inputs
inp = tf.random.uniform([10, 200, 3000])
x = tf.random.uniform([200, 3000])

# Initialize test class
tc = Testclass()

# Process data
output = tf.scan(
    lambda _x, _vin: tc.get_d(_x, _vin),  # Multiple calls to "get_d"
    inp
)

Do the operations in get_c get executed, even though the underlying tensors do not change? If so, what are options for optimizing the code? Is there something better/smarter than putting the tf.scan call inside Testclass?

0 Answers
Related