Is there any problem if a metric state is not a tensor or a list?

[celsofranssa]

Currently, my customized MRR metric combining torch metrics with ranx (a Python library of fast ranking evaluation metrics):

import pickle
import torch
from ranx import Qrels, Run, evaluate
from torchmetrics import Metric

class MRRMetric(Metric):
    def __init__(self, params):
        super(MRRMetric, self).__init__(compute_on_step=False)
        self.params=params
        self.run = {} # validation queries and documents
        self.relevance_map = self._load_relevance_map() # the relevance map


    def _load_relevance_map(self):
        with open(f"{self.params.relevance_map.dir}relevance_map.pkl", "rb") as relevances_file:
            data = pickle.load(relevances_file)
        relevance_map = {}
        for text_idx, labels_ids in data.items():
            d = {}
            for label_idx in labels_ids:
                d[f"label_{label_idx}"] = 1.0
            relevance_map[f"text_{text_idx}"] = d
        return relevance_map

    def similarities(self, x1, x2):
        """
        Calculates the cosine similarity matrix for every pair (i, j),
        where i is an embedding from x1 and j is another embedding from x2.

        :param x1: a tensors with shape [batch_size, hidden_size].
        :param x2: a tensors with shape [batch_size, hidden_size].
        :return: the cosine similarity matrix with shape [batch_size, batch_size].
        """
        x1 = x1 / torch.norm(x1, dim=1, p=2, keepdim=True)
        x2 = x2 / torch.norm(x2, dim=1, p=2, keepdim=True)
        return torch.matmul(x1, x2.t())

    def update(self, query_ids, query_rpr, document_ids, document_rpr):
        similarities = self.similarities(query_rpr, document_rpr)
        for row, query_idx in enumerate(query_ids.tolist()):
            if f"query_{query_idx}" not in self.run:
                self.run[f"query_{query_idx}"]={}
            for col, document_idx in enumerate(document_ids.tolist()):
                self.run[f"query_{query_idx}"][f"document_{document_idx}"] = similarities[row][col].item()

    def compute(self):
        return evaluate(
            Qrels({key: value for key, value in self.relevance_map.items() if key in self.run.keys()}),
            Run(self.run),
            ["mrr"] # could be any ranx metric (https://amenra.github.io/ranx/metrics/)
        )
    
    def reset(self) -> None:
        self.run = {}

Then, ss there any problem using a dictionary as a metric state?