From 7dade54eb4d4d69b98a7aae834b0b1b08e8375a7 Mon Sep 17 00:00:00 2001
From: Bradley Lewis Fargo <blewfargs@gmail.com>
Date: Fri, 20 Mar 2026 01:31:28 -0500
Subject: [PATCH 1/2] Fix Nx.linspace crash with n=1

When n=1 and endpoint=true (default), the divisor is n-1=0, causing
a divide-by-zero. Special-case n=1 to return start value directly.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 nx/lib/nx.ex        | 17 +++++++++++++++++
 nx/test/nx_test.exs |  8 ++++++++
 2 files changed, 25 insertions(+)

diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex
index fc5f55dcc3..c65dc0361e 100644
--- a/nx/lib/nx.ex
+++ b/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
+      start |> new_axis(-1, opts[:name]) |> as_type(opts[:type])
+    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} ->
diff --git a/nx/test/nx_test.exs b/nx/test/nx_test.exs
index 5d42216132..7aa324a769 100644
--- a/nx/test/nx_test.exs
+++ b/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

From 625e071607e533536c098342fd1bdba8ac337c98 Mon Sep 17 00:00:00 2001
From: Paulo Valente <16843419+polvalente@users.noreply.github.com>
Date: Mon, 23 Mar 2026 23:40:15 -0300
Subject: [PATCH 2/2] Apply suggestion from @polvalente

---
 nx/lib/nx.ex | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex
index c65dc0361e..4bffd78c49 100644
--- a/nx/lib/nx.ex
+++ b/nx/lib/nx.ex
@@ -16833,7 +16833,7 @@ defmodule Nx do
 
     if n == 1 do
       # Special case: single point returns start value
-      start |> new_axis(-1, opts[:name]) |> as_type(opts[:type])
+      new_axis(start, -1, opts[:name])
     else
       linspace_n(start, stop, n, opts, vectorized_axes)
     end