[Python] TensorArray read when size is 0

num_items function throws

in user code: File "/Users/anu/PycharmProjects/Siglip/KVCache.py", line 12, in num_items * tf.cond( self.key_cache.size() > 0, OutOfRangeError: Tried to read from index 0 but array size is: 0
​

I couldn't understand why it executes read so eagerly. update isn't debugged yet.

import tensorflow as tf

class KVCache:

def __init__(self, ):
self.key_cache = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
self.value_cache = tf.TensorArray(tf.int32, size=0, dynamic_size=True)


@tf.function
def num_items(self):
tf.cond( self.key_cache.size() > 0,
tf.shape(self.key_cache.read(0))[-2],
tf.constant(0, tf.int32)
)
@tf.function
def update(self,
key_states,
value_states,
layer_idx):
if self.key_cache.size() <= layer_idx:
self.key_cache.write(self.key_cache.size(),key_states)
self.value_cache.write(self.value_cache.size(),value_states)
else:
self.key_cache[layer_idx] = tf.concat([self.key_cache[layer_idx], key_states], axis=-2)
self.value_cache[layer_idx] = tf.concat([self.value_cache[layer_idx], value_states], axis=-2)
return self.key_cache[layer_idx], self.value_cache[layer_idx]

Continue reading...