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?