[Question] can `Scann` be used inside the model during training?

I have the following model using sequences to predict the next item:

class Model(tfrs.models.Model):
    def __init__(self):
        super().__init__()

        self.query_model = tf.keras.Sequential([
            QueryModel(),
            tf.keras.layers.Dense(64),
            L2NormalizationLayer(axis=1)

        ])
        
        self.candidate_model = tf.keras.Sequential([
            CandidateModel(),
            tf.keras.layers.Dense(64),
            L2NormalizationLayer(axis=1)
        ])
        
        scann = tfrs.layers.factorized_top_k.ScaNN(num_reordering_candidates=100)
        scann.index_from_dataset(
            candidates_ds.map(
                lambda x: (x['id'], self.candidate_model({ 'url': x['url'] }))
            )
        )
        
        self.task = tfrs.tasks.Retrieval(
            # Normal approach commented out, using the candidate model to map over a dataset of unique candidates.
             metrics=tfrs.metrics.FactorizedTopK(
                  # candidates=candidates_ds.map( lambda x: (x['id'], self.candidate_model({ 'url': x['url'] })))
                  # Instead we use the SCANN layer
                  candidates=scann
            )
            remove_accidental_hits=True
        )
        
    def call(self, features):
        candidate_embeddings = self.candidate_model({
            'url': features['url'],
        })

        query_embeddings = self.query_model({
            'advertiser_name': features['advertiser_name']        
        })
                
        return (
            query_embeddings,
            candidate_embeddings,
        )
    
    def compute_loss(self, features, training=False):
        query_embeddings, candidate_embeddings = self(features)

        return self.task(
            query_embeddings, 
            candidate_embeddings,
            candidate_ids=features['id'],
            compute_metrics=not training
        )

This runs fine and much quicker! The evaluation step sees 100x speed ups!

However this model does not improve on the metric, my first thought it that the model is not updating the embeddings each epoch as they are run only once. However, in the normal approach we also pass a dataset of already mapped candidate embeddings...

At which point in the training does the model update the embeddings it is learning to use for new evaluation runs?

Trace the ScaNN setup in the model constructor, the candidate_model calls, and compute_loss with compute_metrics. Check when FactorizedTopK and ScaNN receive candidate embeddings during training versus evaluation. Done means documenting whether the index is refreshed and why the metric differs from the normal mapped-dataset approach.