tfr MRR metric for sparse tensors crashes

[fraolBatole]

Hello,

I was trying to use tfr.keras.metrics.MRRMetric() at compile time. But it crashes with a shape-related bug. Below is the code I used:

mrr_metric = tfr.keras.metrics.MRRMetric(topn=10)
def create_classifier_bert_model():    
    inputs = layers.Input((config.MAX_LEN,), dtype=tf.int64)
    sequence_output = pretrained_bert_model(inputs)
    outputs = layers.Dense(vectorize_layer.vocabulary_size(), activation="softmax")(sequence_output)
    # outputs2 = layers.Dense(vectorize_layer.vocabulary_size(), activation="softmax")(sequence_output)
    classifer_model = keras.Model(inputs, outputs=[outputs], name="prediction")
    optimizer = keras.optimizers.Adam(learning_rate=config.LR)
    classifer_model.compile(
        optimizer=optimizer, loss="sparse_categorical_crossentropy", weighted_metrics=[tf.keras.metrics.SparseCategoricalAccuracy(), mrr_metric]
    )
    return classifer_model

classifer_model = create_classifier_bert_model()
print(config.LR)
classifer_model.fit(
    mlm_ds_ft_cs,
    epochs=5,
)

The error is ValueError: Shapes (None, 64, 49934) and (None, 64) are incompatible

I have also tried to reshape (y_true, y_pred) but that causes memory issues. Is there any way to fix this issue, or is there an alternative way to compute MRR in Keras/tensorflow? Thanks!