Update lib/scholar/linear/bayesian_ridge_regression.ex

Co-authored-by: Krsto Proroković <[email protected]>
elixir-nx · Apr 21, 2024 · 4570fc5 · 4570fc5
1 parent 52b2084
commit 4570fc5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lib/scholar/linear/bayesian_ridge_regression.ex b/lib/scholar/linear/bayesian_ridge_regression.ex
@@ -457,7 +457,7 @@ while {{coef, rmse, alpha, lambda, iter = Nx.u64(0), has_converged = Nx.u8(0), s
         {scalar * x, scalar * y}
 
       _ ->
-        scale = sample_weights |> Nx.sqrt() |> Nx.make_diagonal()
+scale = sample_weights |> Nx.sqrt() |> Nx.new_axis(1)
         {Nx.dot(scale, x), Nx.dot(scale, y)}
     end
   end