You could implement the same behavior as scikit (pairwise cosine similarities between vectors in the given matrices) without iteration, by normalizing along the last axis, transposing the second matrix, and performing a matrix multiplication:

defmodule Pairwise do
  import Nx.Defn

  defnp l2_norm(x, opts \\ []) do
    (x * x)
    |> Nx.sum(opts)
    |> Nx.sqrt()
  end

  defnp normalize(x, opts \\ []) do
    x / l2_norm(x, axes: opts[:axes], keep_axes: true)
  end

  defn cosine_similarity(a, b) do
    normalized_a = normalize(a, axes: [-1])
    normalized_b = normalize(b, axes: [-1])

    Nx.dot(normalized_a, Nx.transpose(normalized_b))
  end
end

You can now use it passing two vectors:

a = Nx.tensor([1, 2, 3])
b = Nx.tensor([9, 8, 7])

Pairwise.cosine_similarity(a, b)
# =>
#  #Nx.Tensor<
#    f32
#    0.8826588988304138
#  >

Or passing two matrices, and getting pairwise similarities, like with scikit:

a = Nx.tensor([[1, 2, 3], [4, 5, 6]])
b = Nx.tensor([[9, 8, 7], [6, 5, 4]])

Pairwise.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32[2][2]
#   [
#     [0.8826588988304138, 0.8528028130531311],
#     [0.9654632806777954, 0.948051929473877]
#   ]
>

Pitching in just to say that your function is formally correct, although you could simplify it a bit by using existing Nx functions for the dot product and the norm (possibly benefiting from some optimizations):

defmodule CosSim do
  import Nx.Defn

  defn cosine_similarity(a, b) do
    Nx.dot(a, b) / (Nx.LinAlg.norm(a) * Nx.LinAlg.norm(b))
  end
end

You can verify that it’s correct with some sanity checks:

# Similarity for vectors pointing in the same direction is 1
a = Nx.tensor([1, 2])
b = Nx.tensor([2, 4])
CosSim.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32
#   1.0
# >

# Similarity for orthogonal vectors is 0
a = Nx.tensor([3, 0])
b = Nx.tensor([0, 3])
CosSim.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32
#   0.0
# >

# Similarity for vectors pointing in opposite directions is -1:
a = Nx.tensor([5, 0, 2])
b = Nx.tensor([-10, 0, -4])
CosSim.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32
#   -1.0
# >

# Vectors at 45deg have cosine similarity = cos(Pi/4) ~ 0.7071
a = Nx.tensor([1, 0])
b = Nx.tensor([1, 1])
CosSim.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32
#   0.7071067690849304
# >

Note though that the definition of the cosine distance in scholar also has additional special cases for when the L2 norm is so small that it would cause numerical stability issues, or when one or both the operands have norm equal to zero. Additionally, it allows to calculate the cosine distance between multiple vectors in batch, by passing rank-2 tensors as arguments.

We have Cosine Distance available in Scholar: scholar/lib/scholar/metrics/distance.ex at main · elixir-nx/scholar · GitHub

Cosine Similarity would be 1 - CosDistance