Repository: google/jax

Issue:
Cannot take gradient of function using linspace 
I've found I will get an exception "Tracer can't be used with raw numpy function" when I use linspace where one of the arguments to linspace is a parameter which needs to be differentiated with respect to. Here is a MWE that shows the problem 

```
import jax.numpy as jnp
import jax 

def testlinspace(x, *args): #does not work 
    out = jnp.linspace(x[0],x[1],50) #linspace depends on parameters 
    
    out = x[2]*jnp.square(out)
    out = jnp.sum(out)
    
    return out 

def testlinspace2(x, *args): #works 
    out = jnp.linspace(1.,50.,50) #linspace is just some array of numbers 
    
    out = x[2]*jnp.square(out)+x[0]*x[1]
    out = jnp.sum(out)
    
    return out 

x = [2.,3.,1.5]
testobj = testlinspace(x)
testgrad = jax.grad(testlinspace)
testobj2 = testlinspace2(x)
testgrad2 = jax.grad(testlinspace2)

print(testgrad(x)) #throws exception, note gradient using finite differences is DeviceArray([174.48979293, 200.51020897, 316.83673001])
print(testgrad2(x)) #works, gives DeviceArray([  149.99968698,    99.99930626, 42924.99943404])
```

So the problem is you can't do something like jax.numpy.linspace(x[0],10,50) because the linspace is depending on the parameters 'x' being reverse-mode differentiated.



Discussion:
Yeah, we seem to just be wrapping the original numpy implementation.  The below is a draft of a native differentiable form that mirrors the numpy api close to machine precision for an extensive set of all the linspace arguments that I could think to test.

```python
import numpy as onp
from jax import lax
import jax.numpy as np
def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, 
             axis=0):
  """Implementation of linspace differentiable w.r.t. start and stop args."""
  lax._check_user_dtype_supported(dtype, "linspace")
  dtype = onp.float32 if dtype is None else dtype
  bounds_shape = list(lax.broadcast_shapes(onp.shape(start), onp.shape(stop)))
  axis = len(bounds_shape) if axis == -1 else axis
  bounds_shape.insert(axis, 1)
  iota_shape = [1,] * len(bounds_shape)
  iota_shape[axis] = num
  delta = (stop - start) / num
  if endpoint:
    delta *= num / (num - 1)
  out = (np.reshape(start, bounds_shape) + 
         np.reshape(lax.iota(dtype, num), iota_shape) * 
         np.reshape(delta, bounds_shape))
  if retstep:
    return np.array(out, dtype=dtype), delta
  else:
    return np.array(out, dtype=dtype)
```

you can see it working with your gradient test in the public colab notebook:
https://colab.research.google.com/drive/1bdyUypR-S2stCe2H_LyrEBB4sSFZ9ssL

If this looks good to others we could upstream it to replace the current wrapper implementation.
I can't see the notebook
My apologies, this should be a public link (edited above as well): 
https://colab.research.google.com/drive/1bdyUypR-S2stCe2H_LyrEBB4sSFZ9ssL
Sorry it took me awhile to reply. I looked at your solution and verified it worked for my MWE. I tested this for a more complex example as well and it also worked for me there. So in my opinion this should replace the existing jax linspace.

Please let me know if you want me to close this issue. 

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
@@ -1677,6 +1677,7 @@ def arange(start, stop=None, step=None, dtype=None):
   # Fall back to instantiating an ndarray in host memory
   return onp.arange(start, stop=stop, step=step, dtype=dtype)
 
+
 def _wrap_numpy_nullary_function(f):
   """Adapts `f` to return a DeviceArray instead of an onp.ndarray.
 
@@ -1687,24 +1688,58 @@ def wrapper(*args, **kwargs):
     return asarray(f(*args, **kwargs))
   return wrapper
 
+
+@_wraps(onp.linspace)
 def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
              axis=0):
+  """Implementation of linspace differentiable in start and stop args."""
   lax._check_user_dtype_supported(dtype, "linspace")
-  try:
-    out = onp.linspace(start, stop, num, endpoint, retstep, dtype, axis)
-    if retstep:
-      return asarray(out[0]), out[1]
-    else:
-      return asarray(out)
-  except TypeError:  # Old versions of onp may lack axis arg.
-    out = onp.linspace(start, stop, num, endpoint, retstep, dtype)
-    if retstep:
-      return moveaxis(asarray(out[0]), 0, axis), out[1]
-    else:
-      return moveaxis(asarray(out), 0, axis)
+  dtype = dtype or onp.result_type(start, stop, float(num))
+  bounds_shape = list(lax.broadcast_shapes(shape(start), shape(stop)))
+  broadcast_start = broadcast_to(start, bounds_shape)
+  axis = len(bounds_shape) + axis + 1 if axis < 0 else axis
+  bounds_shape.insert(axis, 1)
+  iota_shape = [1,] * len(bounds_shape)
+  iota_shape[axis] = num
+  if endpoint:
+    delta = (stop - start) / (num - 1)
+  else:
+    delta = (stop - start) / num
+  out = (reshape(broadcast_start, bounds_shape) +
+         reshape(lax.iota(dtype, num), iota_shape) *
+         reshape(delta, bounds_shape))
+  if retstep:
+    return lax.convert_element_type(out, dtype), delta
+  else:
+    return lax.convert_element_type(out, dtype)
+
+
+@_wraps(onp.logspace)
+def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
+  """Implementation of logspace differentiable in start and stop args."""
+  lin = linspace(start, stop, num,
+                 endpoint=endpoint, retstep=False, dtype=None, axis=axis)
+  if dtype is None:
+    return power(base, lin)
+  else:
+    return lax.convert_element_type(power(base, lin), dtype)
+
+
+@_wraps(onp.geomspace)
+def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
+  """Implementation of geomspace differentiable in start and stop args."""
+  dtype = dtype or onp.result_type(start, stop, float(num),
+                                   zeros((), dtype))
+  # follow the numpy geomspace convention for negative and complex endpoints
+  signflip = 1 - (1 - sign(real(start))) * (1 - sign(real(stop))) // 2
+  res = signflip * logspace(log10(signflip * start),
+                            log10(signflip * stop), num,
+                            endpoint=endpoint, base=10.0,
+                            dtype=dtype, axis=0)
+  if axis != 0:
+    res = moveaxis(res, 0, axis)
+  return lax.convert_element_type(res, dtype)
 
-logspace = _wrap_numpy_nullary_function(onp.logspace)
-geomspace = _wrap_numpy_nullary_function(onp.geomspace)
 
 @_wraps(onp.meshgrid)
 def meshgrid(*args, **kwargs):


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
@@ -2034,6 +2034,167 @@ def testMeshGrid(self, shapes, dtype, indexing, sparse, rng_factory):
     lnp_fun = partial(lnp.meshgrid, indexing=indexing, sparse=sparse)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
+  def assertDowncastDtypeEqual(self, x, y):
+    """Heuristic for comparing numpy and jax downcast dtypes."""
+    x_dt = jtu._dtype(x)
+    y_dt = jtu._dtype(y)
+    testing_tpu = jtu.device_under_test().startswith("tpu")
+    testing_x32 = not jax.config.read('jax_enable_x64')
+    to32dtype = {onp.int64: onp.int32, onp.uint64: onp.uint32,
+                 onp.float64: onp.float32, onp.float128: onp.float32,
+                 onp.complex128: onp.complex64, onp.complex256: onp.complex64}
+    to32dtype = {onp.dtype(k): onp.dtype(v) for k,v in to32dtype.items()}
+    if testing_tpu or testing_x32:
+      x_dt = to32dtype.get(x_dt, x_dt)
+      y_dt = to32dtype.get(y_dt, y_dt)
+    assert x_dt == y_dt, "truncated dtypes %s != %s" % (x_dt, y_dt)
+
+  @parameterized.named_parameters(
+      jtu.cases_from_list(
+        {"testcase_name": ("_start_shape={}_stop_shape={}_num={}_endpoint={}"
+                           "_retstep={}_dtype={}").format(
+            start_shape, stop_shape, num, endpoint, retstep, dtype),
+         "start_shape": start_shape, "stop_shape": stop_shape,
+         "num": num, "endpoint": endpoint, "retstep": retstep,
+         "dtype": dtype, "rng_factory": rng_factory}
+        for start_shape in [(), (2,), (2, 2)]
+        for stop_shape in [(), (2,), (2, 2)]
+        for num in [5, 20]
+        for endpoint in [True, False]
+        for retstep in [True, False]
+        for dtype in number_dtypes + [None,]
+        for rng_factory in [jtu.rand_default]))
+  def testLinspace(self, start_shape, stop_shape, num, endpoint,
+                   retstep, dtype, rng_factory):
+    rng = rng_factory()
+    # relax default tolerances slightly
+    tol = tolerance(dtype if dtype else onp.float32) * 10
+    args_maker = self._GetArgsMaker(rng,
+                                    [start_shape, stop_shape],
+                                    [dtype, dtype])
+    start, stop = args_maker()
+    ndim = len(onp.shape(start + stop))
+    for axis in range(-ndim, ndim):
+      lnp_op = lambda start, stop: lnp.linspace(
+        start, stop, num,
+        endpoint=endpoint, retstep=retstep, dtype=dtype, axis=axis)
+      onp_op = lambda start, stop: onp.linspace(
+        start, stop, num,
+        endpoint=endpoint, retstep=retstep, dtype=dtype, axis=axis)
+      self._CheckAgainstNumpy(onp_op, lnp_op, args_maker,
+                              check_dtypes=False, tol=tol)
+      # Check dtype equivalence within expected 32bit downcasting.
+      a, b = lnp_op(start, stop), onp_op(start, stop)
+      if retstep:
+        self.assertDowncastDtypeEqual(a[0], b[0])
+        self.assertDowncastDtypeEqual(a[1], b[1])
+      else:
+        self.assertDowncastDtypeEqual(a, b)
+      # floating-point compute between jitted platforms and non-jit + rounding
+      # cause unavoidable variation in integer truncation for some inputs.
+      if dtype in (inexact_dtypes + [None,]):
+        self._CompileAndCheck(lnp_op, args_maker,
+                              check_dtypes=False, atol=tol, rtol=tol)
+
+  @parameterized.named_parameters(
+      jtu.cases_from_list(
+        {"testcase_name": ("_start_shape={}_stop_shape={}_num={}_endpoint={}"
+                           "_base={}_dtype={}").format(
+            start_shape, stop_shape, num, endpoint, base, dtype),
+         "start_shape": start_shape,
+         "stop_shape": stop_shape,
+         "num": num, "endpoint": endpoint, "base": base,
+         "dtype": dtype, "rng_factory": rng_factory}
+        for start_shape in [(), (2,), (2, 2)]
+        for stop_shape in [(), (2,), (2, 2)]
+        for num in [5, 20]
+        for endpoint in [True, False]
+        for base in [10.0, 2, onp.e]
+        for dtype in number_dtypes + [None,]
+        for rng_factory in [jtu.rand_default]))
+  def testLogspace(self, start_shape, stop_shape, num,
+                   endpoint, base, dtype, rng_factory):
+    if (dtype in int_dtypes and
+        jtu.device_under_test() == "gpu" and
+        not FLAGS.jax_enable_x64):
+      raise unittest.SkipTest("GPUx32 truncated exponentiation"
+                              " doesn't exactly match other platforms.")
+    rng = rng_factory()
+    # relax default tolerances slightly
+    tol = tolerance(dtype if dtype else onp.float32) * 10
+    args_maker = self._GetArgsMaker(rng,
+                                    [start_shape, stop_shape],
+                                    [dtype, dtype])
+    start, stop = args_maker()
+    ndim = len(onp.shape(start + stop))
+    for axis in range(-ndim, ndim):
+      lnp_op = lambda start, stop: lnp.logspace(
+        start, stop, num, endpoint=endpoint, base=base, dtype=dtype, axis=axis)
+      onp_op = lambda start, stop: onp.logspace(
+        start, stop, num, endpoint=endpoint, base=base, dtype=dtype, axis=axis)
+      self._CheckAgainstNumpy(onp_op, lnp_op, args_maker,
+                              check_dtypes=False, tol=tol)
+      # Check dtype equivalence within expected 32bit downcasting.
+      a, b = lnp_op(start, stop), onp_op(start, stop)
+      self.assertDowncastDtypeEqual(a, b)
+      if dtype in (inexact_dtypes + [None,]):
+        # Why do compiled and op-by-op float16 np.power numbers differ
+        # slightly more than expected?
+        atol = tol if dtype != onp.float16 else 10 * tol
+        self._CompileAndCheck(lnp_op, args_maker,
+                              check_dtypes=False, atol=atol, rtol=tol)
+
+  @parameterized.named_parameters(
+      jtu.cases_from_list(
+        {"testcase_name": ("_start_shape={}_stop_shape={}_num={}_endpoint={}"
+                           "_dtype={}").format(
+            start_shape, stop_shape, num, endpoint, dtype),
+         "start_shape": start_shape,
+         "stop_shape": stop_shape,
+         "num": num, "endpoint": endpoint,
+         "dtype": dtype, "rng_factory": rng_factory}
+        for start_shape in [(), (2,), (2, 2)]
+        for stop_shape in [(), (2,), (2, 2)]
+        for num in [5, 20]
+        for endpoint in [True, False]
+        # NB: numpy's geomspace gives nonsense results on integer types
+        for dtype in inexact_dtypes + [None,]
+        for rng_factory in [jtu.rand_default]))
+  def testGeomspace(self, start_shape, stop_shape, num,
+                    endpoint, dtype, rng_factory):
+    rng = rng_factory()
+    # relax default tolerances slightly
+    tol = tolerance(dtype if dtype else onp.float32) * 10
+    def args_maker():
+      """Test the set of inputs onp.geomspace is well-defined on."""
+      start, stop = self._GetArgsMaker(rng,
+                                [start_shape, stop_shape],
+                                [dtype, dtype])()
+      # onp.geomspace can't handle differently ranked tensors
+      # w. negative numbers!
+      start, stop = lnp.broadcast_arrays(start, stop)
+      if dtype in complex_dtypes:
+        return start, stop
+      # to avoid NaNs, non-complex start and stop cannot
+      # differ in sign, elementwise
+      start = start * lnp.sign(start) * lnp.sign(stop)
+      return start, stop
+    start, stop = args_maker()
+    ndim = len(onp.shape(start + stop))
+    for axis in range(-ndim, ndim):
+      lnp_op = lambda start, stop: lnp.geomspace(
+        start, stop, num, endpoint=endpoint, dtype=dtype, axis=axis)
+      onp_op = lambda start, stop: onp.geomspace(
+        start, stop, num, endpoint=endpoint, dtype=dtype, axis=axis)
+      self._CheckAgainstNumpy(onp_op, lnp_op, args_maker,
+                              check_dtypes=False, tol=tol)
+      # Check dtype equivalence within expected 32bit downcasting.
+      a, b = lnp_op(start, stop), onp_op(start, stop)
+      self.assertDowncastDtypeEqual(a, b)
+      if dtype in (inexact_dtypes + [None,]):
+        self._CompileAndCheck(lnp_op, args_maker,
+                              check_dtypes=False, atol=tol, rtol=tol)
+
   def testDisableNumpyRankPromotionBroadcasting(self):
     try:
       prev_flag = FLAGS.jax_numpy_rank_promotion