How to get Token.similarity using spacy 3.7+ ?

[wildfluss]

Kinda follow-up on this #13137 (comment) but I'm just trying to get token similarity using en_core_web_trf and latest spacy

This

import spacy
nlp = spacy.load("en_core_web_trf")

doc_city_trf = nlp("Helsinki is the capital of Finland.")
# doc_money_trf = nlp("Filand has the capital of Helsinki.")
# city_trf = doc_city_trf[3]
# money_trf = doc_money_trf[3]
# sim = city_trf.similarity(money_trf)

for token in doc_city_trf:
    print(token.text, token.similarity(doc_city_trf))

is trying Evaluating Token.similarity based on empty vectors

and then trying to get cosine similarity like this

a = doc_memo._.trf_data.last_hidden_layer_state[t_memo.i].data
b = doc_some_eq._.trf_data.last_hidden_layer_state[t_some_eq.i].data
        
sim[i, j] = dot(a, b)/(norm(a)*norm(b))

for some tokens from two documents kinda does not work because sometimes

> (doc_some_eq._.trf_data.last_hidden_layer_state[2].data).shape
(4, 768)

and for other token indecies its (1, 768) (which I would expect) I guess that has to do with "word pieces" or something but would it be correct to just average (?, 768) to get comparable tensors for a rough word vector similarity ?