Repository: google/jax

Issue:
`jnp.nonzero` zero padding behavior is confusing
I noticed that `jax.numpy.nonzero` now allows the wonderful `size` parameter which allows jitting: #6501. This is wonderful! Thanks!

However, the semantics is that missing zeros are zero-padded. This seems unworkable in some situations, consider:
```python
x = jnp.zeros(5)
print(jnp.nonzero(x, size=1))   # Prints [ 0 ]

x = jnp.ones(5)
print(jnp.nonzero(x, size=1))   # Prints [ 0 ]
```

There's no way to tell between these two cases.

Instead, perhaps you could consider something like `-1`-padding? 
Though, performance-wise, I'm not sure how hard it would be to make that on-par with the current implementation 


Discussion:
 Note that this only occurs when `size=1`. 
 
So perhaps at least we should document this corner case since this seems like a problem that could cause unexpected program behaviors.
Thanks for the question – we chose to make the padding zero because it's the natural value that falls out of the algorithm used to implement the function. Zero padding makes sense in many applications, and any other fill value would require additional post-processing of the result. If you wish, you can do this post-processing yourself. For example:
```python
import jax.numpy as jnp

def nonzero_with_configurable_padding(arr, *, size=None, fill_value=-1):
  out = jnp.nonzero(arr, size=size)
  if size is not None:
    num_nonzero = (arr != 0).sum()
    mask = jnp.arange(len(out[0])) < num_nonzero
    out = tuple(jnp.where(mask, entry, fill_value) for entry in out)
  return out

dx = jnp.zeros(5)
print(nonzero_with_configurable_padding(x, size=1))
# [ -1 ]

x = jnp.ones(5)
print(nonzero_with_configurable_padding(x, size=1))
# [ 0 ]
```
@jakevdp Thanks for the explanation and your example. 
A couple questions for this overall issue:
1) Is there an alternative algorithm that can a) use arbitrary `fill_value` b) retain the same efficiency?
2) Should we document this behavior for `size=1`? I would say that we could even ban `size=1` since it produces ambiguous return value.
3) Perhaps you can add this implementation of `fill_value` parameter to `nonzero`? The default is zero so the fast code path is enabled. 

Ideally, someone will think of solution to 1), but before that, I think we should do either 2) or 3). 
WDYT? 

I think the behavior for size=1 is already documented, no? The true size is zero, so the result is padded with zeros up to the requested size.
> The true size is zero, so the result is padded with zeros up to the requested size.

Well, sure. I'm simply thinking that we should add a note, that `size=1` will have this problem of ambiguous output. 
I mean, from a user's perspective, I do hope that when I write some code with `size=1`, I can get fair amount of warning of this behavior. Maybe it's just me, but I only realized this issue after I observed unexpected output. 

I can do a PR on the docs. 
The only reason I haven't done so is because I'm not sure if 1) or 3) is a viable/better option?
I think a `fill_value` option may be a reasonable feature to add. The reason I didn't before is that it adds a maintenance/testing burden for something that didn't have any concrete use-case to motivate it, and the relevant information (the true number of nonzeros) is straightforward to determine separately.

Can you give an example of code you might write with a configurable fill value?
Sure. The code I had which ran into this problem was a jittable [Schur-Transform](https://en.wikipedia.org/wiki/Lehmer–Schur_algorithm). The implementation is pretty short so I'd paste it here (note that most of it is just comments). 

```python
import jax.numpy as jnp


def _first_non_zero(x):
  return jnp.nonzero(x, size=1)[0][0]
  # return jnp.nonzero(x, size=1, fill_value=x.size)[0][0]


def schur_transform(coeffs,
                    n_leading_zeros: int):
  """Schur Transform of a polynomial.
  
  Computes 
  ..math::

    T(p) = \\bar{p(0)}p - \\bar{p^*(0)}p^*
  
  where :math:`p^*(z)` is the reciprocal adjoint polynomial 
  ..math::

    p^*(z) = z^n \cdot \\bar{p(z^{-1})}

  and bar denotes complex conjugation.

  See https://en.wikipedia.org/wiki/Lehmer–Schur_algorithm for more details.


  Args:
    coeffs: a complex array of polynomial coefficients, with possibily some 
    leading zero terms.
    n_leading_zeros: the number of leading zeros in the polynomial.

  Returns:
    the transformed polynomial coefficients, and the number of leading zeros in
    it.
  """
  # The reciprocal adjoint polynomial
  coeffs_ra = jnp.roll(coeffs.conj(), shift=-n_leading_zeros, axis=0)[::-1]
  # Schur Transformation
  coeffs = (coeffs[-1].conj() * coeffs - coeffs_ra[-1].conj() * coeffs_ra)
  # Update number of leading zeros
  n_leading_zeros = _first_non_zero(coeffs)
  return coeffs, n_leading_zeros
```

The test could be
```
coeffs = jnp.array([1, 2, 1], dtype=jnp.complex64) # f(x) = x^2 + 2x + 1
transformed_coeffs, n_leading_zeros = schur_transform(coeffs, n_leading_zeros=0)
assert jnp.all(transformed_coeffs == 0)
assert n_leading_zeros == 3 # Currently fails because `fill_value` is not available.
```

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2321,11 +2321,12 @@ def count_nonzero(a, axis: Optional[Union[int, Tuple[int, ...]]] = None,
 typically compatible with JIT. The JAX version adds the optional `size` argument which
 specifies the size of the output arrays: it must be specified statically for ``jnp.nonzero``
 to be traced. If specified, the first `size` nonzero elements will be returned; if there
-are fewer nonzero elements than `size` indicates, the index arrays will be zero-padded.
+are fewer nonzero elements than `size` indicates, the result will be padded with ``fill_value``,
+which defaults to zero.
 """
 
 @_wraps(np.nonzero, lax_description=_NONZERO_DOC)
-def nonzero(a, *, size=None):
+def nonzero(a, *, size=None, fill_value=None):
   a = atleast_1d(a)
   mask = a != 0
   if size is None:
@@ -2337,7 +2338,13 @@ def nonzero(a, *, size=None):
     return tuple(zeros(size, int) for dim in a.shape)
   flat_indices = cumsum(bincount(cumsum(mask), length=size))
   strides = np.cumprod(a.shape[::-1])[::-1] // a.shape
-  return tuple((flat_indices // stride) % size for stride, size in zip(strides, a.shape))
+  out = tuple((flat_indices // stride) % size for stride, size in zip(strides, a.shape))
+  if size is not None and fill_value is not None:
+    if ndim(fill_value) != 0:
+      raise ValueError(f"fill_value must be a scalar; got {fill_value}")
+    fill_mask = arange(size) >= mask.sum()
+    out = tuple(where(fill_mask, fill_value, entry) for entry in out)
+  return out
 
 @_wraps(np.flatnonzero, lax_description=_NONZERO_DOC)
 def flatnonzero(a, *, size=None):


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
@@ -940,13 +940,14 @@ def testNonzero(self, shape, dtype):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_shape={}_size={}".format(
-          jtu.format_shape_dtype_string(shape, dtype), size),
-       "shape": shape, "dtype": dtype, "size": size}
+      {"testcase_name": "_shape={}_size={}_fill_value={}".format(
+          jtu.format_shape_dtype_string(shape, dtype), size, fill_value),
+       "shape": shape, "dtype": dtype, "size": size, "fill_value": fill_value}
       for shape in nonempty_array_shapes
       for dtype in all_dtypes
+      for fill_value in [None, -1]
       for size in [1, 5, 10]))
-  def testNonzeroSize(self, shape, dtype, size):
+  def testNonzeroSize(self, shape, dtype, size, fill_value):
     rng = jtu.rand_some_zero(self.rng())
     args_maker = lambda: [rng(shape, dtype)]
     @jtu.ignore_warning(category=DeprecationWarning, message="Calling nonzero on 0d arrays.*")
@@ -955,9 +956,9 @@ def np_fun(x):
       if size <= len(result[0]):
         return tuple(arg[:size] for arg in result)
       else:
-        return tuple(np.concatenate([arg, np.zeros(size - len(arg), arg.dtype)])
+        return tuple(np.concatenate([arg, np.full(size - len(arg), fill_value or 0, arg.dtype)])
                      for arg in result)
-    jnp_fun = lambda x: jnp.nonzero(x, size=size)
+    jnp_fun = lambda x: jnp.nonzero(x, size=size, fill_value=fill_value)
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
     self._CompileAndCheck(jnp_fun, args_maker)