Fix type system reports

Author: josevalim

lib/scholar/interpolation/cubic_spline.ex

@@ -241,15 +241,7 @@ defmodule Scholar.Interpolation.CubicSpline do
 
   defnp predict_n(%__MODULE__{x: x, coefficients: coefficients}, target_x, opts) do
     original_shape = Nx.shape(target_x)
-
-    target_x =
-      case target_x do
-        {} ->
-          Nx.new_axis(target_x, 0)
-
-        _ ->
-          Nx.flatten(target_x)
-      end
+    target_x = Nx.flatten(target_x)
 
     idx_selector = Nx.new_axis(target_x, 1) > Nx.new_axis(x, 0)
 

lib/scholar/shared.ex

@@ -55,12 +55,12 @@ defmodule Scholar.Shared do
       is_nil(weights) ->
         Nx.tensor(1.0, type: type)
 
-      Nx.is_tensor(weights) and Nx.shape(weights) in [{}, {num_samples}] ->
-        weights |> Nx.broadcast({num_samples}) |> Nx.as_type(type)
-
       is_list(weights) and length(weights) == num_samples ->
         Nx.tensor(weights, type: type)
 
+      Nx.is_tensor(weights) and Nx.shape(weights) in [{}, {num_samples}] ->
+        weights |> Nx.broadcast({num_samples}) |> Nx.as_type(type)
+
       true ->
         raise ArgumentError,
               "invalid value for :weights option: expected list or tensor of positive numbers of size #{num_samples}, got: #{inspect(weights)}"

test/scholar/interpolation/cubic_spline_test.exs

@@ -206,6 +206,15 @@ defmodule Scholar.Interpolation.CubicSplineTest do
       end
     end
 
+    test "predict/2 supports scalar tensors with jit_apply" do
+      model = CubicSpline.fit(Nx.iota({3}), Nx.iota({3}))
+
+      prediction =
+        Nx.Defn.jit_apply(&CubicSpline.predict/3, [model, Nx.tensor(1.0), []])
+
+      assert prediction == Nx.tensor(1.0)
+    end
+
     test "not sorted x" do
       x = Nx.tensor([3, 2, 4, 1, 0])
       y = Nx.tensor([-10, 3, -1, 2, 1])