How to do Cosine similarity between 2 matrices with same concept as of scikit-learn?

I have one tensor x of shape [64,324]. Think of 64 as batch size and 324 as number of features. I want to perform cosine similarity with another tensor w such that that the output tensor is of size [64,10] , where 64 is same as batch size and 10 is output features. I could not find any tensor w which serves this purpose. w should be independent of batch size.

For x of shape 64x324, if w is 324x10 the product x @ w is 64x10 and is 64 x 10 scalar products (so higher = closer alignment).
You effectively compare the batch features against 10 example features and measure alignment to each of them.


Thanks very much for your reply