Tensor Comprehensions and einpos how to use for made a rolling metrics between two vectors

1

I need to implement a cosine similarity between a smallest 1D tensor A and a longer 1D tensor B, I want to iterate the B tensor like a 1D conv and make cosine distance like the following pseudo code:

cosine_sim = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
def cosine_rolling_1(a, b):
    for i in range(len(b) - len(a) + 1):
        yield cosine_sim(a, b[i:i+len(a)])

def cosine_rolling_2(a, b):
    for i in range(len(b) - len(a) + 1):
        yield torch.dot(a, b[i:i+len(a)])/(torch.norm(a)*torch.norm(b[i:i+len(a)]))

I thin that einpos is not food because I have two input but also with Tensor Comprehension i dont think that I can write something like that:

def cosine_rolling_tc(float(R) A, float(C) b) -> (o) {
    o(r) +=! dot(A, b(r-c:r-c+R))/(norm(A)*norm(b(r-c:r-c+R)))
}

What is the best way to implement in torch that operation? Do I need to use low level cuda?