xla

Adds torch.nn.functional.cosine_similarity Jax implementation

#8305

Merged