Repository: google/jax

Issue:
jax.numpy.linspace() does not faithfully recover endpoints
Example:
```python
>>> import numpy as np
>>> import jax.numpy as jnp
>>> a, b = -2.3328583, 2.8459014
>>> np.linspace(a, b, 10)[-1] == b
True
>>> jnp.linspace(a, b, 10)[-1] == b
DeviceArray(False, dtype=bool)
```


Discussion:
Looks to be due to float32 roundoff error:
```
>>> b - a, float32(b) - float32(a)
(5.1787598, 5.1787596)
```
...and numpy side-steps this by doing all computation in 64-bit, even if input and output dtypes are 32-bit: https://github.com/numpy/numpy/blob/fc1f196584ff6dd530982febba2322679317c632/numpy/core/function_base.py#L120-L121
I think we could address this by factoring out operations to prevent loss of precision; e.g.
```python
def linspace_bad(a, b, N):
  a, b = np.float32(a), np.float32(b)
  return a + (b - a) * np.arange(N, dtype=np.float32) / (N - 1)

def linspace_good(a, b, N):
  a, b = np.float32(a), np.float32(b)
  step = np.arange(N, dtype=np.float32) / (N - 1)
  return a * (1 - step) + b * step

print(linspace_bad(-2.3328583, 2.8459015, 2))
print(linspace_good(-2.3328583, 2.8459015, 2))
```
```
[-2.3328583  2.8459013]
[-2.3328583  2.8459015]
```

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2160,6 +2160,7 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
   dtype = dtype or dt
   bounds_shape = list(lax.broadcast_shapes(shape(start), shape(stop)))
   broadcast_start = broadcast_to(start, bounds_shape)
+  broadcast_stop = broadcast_to(stop, bounds_shape)
   axis = len(bounds_shape) + axis + 1 if axis < 0 else axis
   bounds_shape.insert(axis, 1)
   iota_shape = [1,] * len(bounds_shape)
@@ -2167,9 +2168,18 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
   div = (num - 1) if endpoint else num
   if num > 1:
     delta = lax.convert_element_type(stop - start, dt) / div
-    out = (reshape(broadcast_start, bounds_shape) +
-           reshape(lax.iota(dt, num), iota_shape) *
-           reshape(delta, bounds_shape))
+    if issubdtype(dtype, integer):
+      # This is similar to how numpy computes linspace, but it
+      # can fail to recover the endpoints in float32 arithmetic.
+      out = (reshape(broadcast_start, bounds_shape) +
+        reshape(lax.iota(dt, num), iota_shape) *
+        reshape(delta, bounds_shape))
+    else:
+      # This approach recovers the endpoints with float32 arithmetic,
+      # but can lead to rounding errors for integer outputs.
+      step = reshape(lax.iota(dt, num), iota_shape) / div
+      out = (reshape(broadcast_start, bounds_shape) * (1 - step) +
+        reshape(broadcast_stop, bounds_shape) * step)
   elif num == 1:
     delta = nan if endpoint else lax.convert_element_type(stop - start, dt)
     out = reshape(broadcast_start, bounds_shape)


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2912,6 +2912,20 @@ def testLinspace(self, start_shape, stop_shape, num, endpoint,
         self._CompileAndCheck(jnp_op, args_maker,
                               check_dtypes=False, atol=tol, rtol=tol)
 
+  @parameterized.named_parameters(
+      jtu.cases_from_list(
+        {"testcase_name": "_dtype={}".format(dtype),
+         "dtype": dtype,
+         "rng_factory": rng_factory}
+        for dtype in number_dtypes
+        for rng_factory in [jtu.rand_default]))
+  def testLinspaceEndpoints(self, dtype, rng_factory):
+    """Regression test for Issue #3014."""
+    rng = rng_factory(self.rng())
+    endpoints = rng((2,), dtype)
+    out = jnp.linspace(*endpoints, 10, dtype=dtype)
+    self.assertAllClose(out[[0, -1]], endpoints, check_dtypes=True, rtol=0, atol=0)
+
   @parameterized.named_parameters(
       jtu.cases_from_list(
         {"testcase_name": ("_start_shape={}_stop_shape={}_num={}_endpoint={}"