Skip to content

Commit 7dade54

Browse files
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>
1 parent 11c6cc5 commit 7dade54

2 files changed

Lines changed: 25 additions & 0 deletions

File tree

nx/lib/nx.ex

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16800,6 +16800,14 @@ defmodule Nx do
1680016800
]
1680116801
>
1680216802
16803+
Linspace with a single point returns the start value:
16804+
16805+
iex> Nx.linspace(0, 10, n: 1)
16806+
#Nx.Tensor<
16807+
f32[1]
16808+
[0.0]
16809+
>
16810+
1680316811
## Error cases
1680416812
1680516813
iex> Nx.linspace(0, 24, n: 1.0)
@@ -16823,6 +16831,15 @@ defmodule Nx do
1682316831
raise ArgumentError, "expected n to be a non-negative integer, got: #{inspect(n)}"
1682416832
end
1682516833

16834+
if n == 1 do
16835+
# Special case: single point returns start value
16836+
start |> new_axis(-1, opts[:name]) |> as_type(opts[:type])
16837+
else
16838+
linspace_n(start, stop, n, opts, vectorized_axes)
16839+
end
16840+
end
16841+
16842+
defp linspace_n(start, stop, n, opts, vectorized_axes) do
1682616843
{iota_shape, start, stop} =
1682716844
case {start.shape, stop.shape} do
1682816845
{shape, shape} ->

nx/test/nx_test.exs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3071,6 +3071,14 @@ defmodule NxTest do
30713071
expected_linear = Nx.tensor([0.0, 0.1, 0.2, 0.3, 0.4, 0.5], type: :f64)
30723072
assert_all_close(linear, expected_linear, atol: 1.0e-15, rtol: 1.0e-15)
30733073
end
3074+
3075+
test "n=1 returns start value" do
3076+
assert Nx.linspace(0, 10, n: 1) == Nx.tensor([0.0])
3077+
end
3078+
3079+
test "n=1 with same start/stop" do
3080+
assert Nx.linspace(5, 5, n: 1) == Nx.tensor([5.0])
3081+
end
30743082
end
30753083

30763084
describe "reflect/2" do

0 commit comments

Comments
 (0)