Fix Nx.linspace crash with n=1, NaN as well

nx/lib/nx.ex

@@ -16800,6 +16800,14 @@ defmodule Nx do
         ]
       >
 
+  Linspace with a single point returns the start value:
+
+      iex> Nx.linspace(0, 10, n: 1)
+      #Nx.Tensor<
+        f32[1]
+        [0.0]
+      >
+
   ## Error cases
 
       iex> Nx.linspace(0, 24, n: 1.0)
@@ -16823,6 +16831,15 @@ defmodule Nx do
       raise ArgumentError, "expected n to be a non-negative integer, got: #{inspect(n)}"
     end
 
+    if n == 1 do
+      # Special case: single point returns start value
+      new_axis(start, -1, opts[:name])
+    else
+      linspace_n(start, stop, n, opts, vectorized_axes)
+    end
+  end
+
+  defp linspace_n(start, stop, n, opts, vectorized_axes) do
     {iota_shape, start, stop} =
       case {start.shape, stop.shape} do
         {shape, shape} ->

nx/test/nx_test.exs

@@ -3071,6 +3071,14 @@ defmodule NxTest do
       expected_linear = Nx.tensor([0.0, 0.1, 0.2, 0.3, 0.4, 0.5], type: :f64)
       assert_all_close(linear, expected_linear, atol: 1.0e-15, rtol: 1.0e-15)
     end
+
+    test "n=1 returns start value" do
+      assert Nx.linspace(0, 10, n: 1) == Nx.tensor([0.0])
+    end
+
+    test "n=1 with same start/stop" do
+      assert Nx.linspace(5, 5, n: 1) == Nx.tensor([5.0])
+    end
   end
 
   describe "reflect/2" do