diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex index fc5f55dcc3..4bffd78c49 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 + 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} -> 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