xla

c9e02e9d - Adds torch.nn.functional.cosine_similarity Jax implementation (#8305)