From 2a28cd4cafddd04ae9b98d6b4cd03f77a27382ff Mon Sep 17 00:00:00 2001
From: srzeszut <szczepan@fedora.com>
Date: Wed, 27 Nov 2024 14:03:28 +0100
Subject: [PATCH] apply suggestions

---
 lib/scholar/naive_bayes/bernoulli.ex | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/scholar/naive_bayes/bernoulli.ex b/lib/scholar/naive_bayes/bernoulli.ex
index 40a0e9a4..eba3cdc4 100644
--- 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