apply suggestions

elixir-nx · Nov 27, 2024 · 2a28cd4 · 2a28cd4
1 parent 572e6c6
commit 2a28cd4
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/lib/scholar/naive_bayes/bernoulli.ex b/lib/scholar/naive_bayes/bernoulli.ex
@@ -355,8 +355,7 @@ defmodule Scholar.NaiveBayes.Bernoulli do
     log_proba_x =
       jll
       |> Nx.logsumexp(axes: [1])
-      |> Nx.new_axis(1)
-      |> Nx.broadcast(jll)
+      |> Nx.reshape({Nx.axis_size(jll, 0), 1})
 
     jll - log_proba_x
   end