Repository: google/jax

Issue:
Jax options not threaded through to base implementation in jnp
**Bug**
`jnp.flatnonzero` and `jnp.argwhere` mention `fill_value` option in documentation (see image below), but calling with `fill_value` leads to 'unexpected keyword argument' error.

In the below code, the last two lines will raise errors due to unexpected argument `fill_value`
```python
import jax.numpy as jnp
d = jnp.array([True,False,True])
jnp.nonzero(d, size=20, fill_value=100)
jnp.flatnonzero(d, size=20, fill_value=100)
jnp.argwhere(d, size=20, fill_value=100)
```

**Why?**
The implementation of `jnp.flatnonzero` and `jnp.argwhere` link back to `jnp.nonzero` but the `fill_value` argument wasn't threaded through (just the `size` argument).

The threading is there in `jnp.where` and other cases, just those two have been missed as far as I've noticed.

**Fix**
Pull request:
#8926 
just threads the argument

Images

![image](https://user-images.githubusercontent.com/44170519/145838724-958be84b-68b3-4d8d-980e-4dc8bc5bb0d3.png)

![image](https://user-images.githubusercontent.com/44170519/145838907-673fb312-8b0b-4433-9be3-0d29cdf5cd82.png)

![image](https://user-images.githubusercontent.com/44170519/145841134-5735c73d-ef6f-4651-8e5e-d08f149a0b81.png)



Discussion:

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
@@ -2875,8 +2875,8 @@ def nonzero(a, *, size=None, fill_value=None):
   return out
 
 @_wraps(np.flatnonzero, lax_description=_NONZERO_DOC)
-def flatnonzero(a, *, size=None):
-  return nonzero(ravel(a), size=size)[0]
+def flatnonzero(a, *, size=None, fill_value=None):
+  return nonzero(ravel(a), size=size, fill_value=fill_value)[0]
 
 
 def _nan_reduction(a, name, jnp_reduction, init_val, nan_if_all_nan,
@@ -5212,8 +5212,8 @@ def vander(x, N=None, increasing=False):
 """
 
 @_wraps(np.argwhere, lax_description=_ARGWHERE_DOC)
-def argwhere(a, *, size=None):
-  result = transpose(vstack(nonzero(a, size=size)))
+def argwhere(a, *, size=None, fill_value=None):
+  result = transpose(vstack(nonzero(a, size=size, fill_value=fill_value)))
   if ndim(a) == 0:
     return result[:0].reshape(result.shape[0], 0)
   return result.reshape(result.shape[0], ndim(a))


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
@@ -998,6 +998,29 @@ def testFlatNonzero(self, shape, dtype):
     jnp_fun = lambda x: jnp.flatnonzero(x, size=np.size(x) // 2)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"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, 10, (-1,), (10,)]
+      for size in [1, 5, 10]))
+  def testFlatNonzeroSize(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.*")
+    def np_fun(x):
+      result = np.flatnonzero(x)
+      if size <= len(result):
+        return result[:size]
+      else:
+        fill_val = fill_value or 0
+        return np.concatenate([result, np.full(size - len(result), fill_val, result.dtype)])
+    jnp_fun = lambda x: jnp.flatnonzero(x, size=size, fill_value=fill_value)
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
+    self._CompileAndCheck(jnp_fun, args_maker)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}".format(
           jtu.format_shape_dtype_string(shape, dtype)),
@@ -1017,6 +1040,30 @@ def testArgWhere(self, shape, dtype):
     jnp_fun = lambda x: jnp.argwhere(x, size=np.size(x) // 2)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"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, shape or (1,)]
+      for size in [1, 5, 10]))
+  def testArgWhereSize(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.*")
+    def np_fun(x):
+      result = np.argwhere(x)
+      if size <= len(result):
+        return result[:size]
+      else:
+        fillvals = fill_value if np.ndim(fill_value) else result.shape[-1] * [fill_value or 0]
+        return np.empty((size, 0), dtype=int) if np.ndim(x) == 0 else np.stack([np.concatenate([arg, np.full(size - len(arg), fval, arg.dtype)])
+                        for fval, arg in safe_zip(fillvals, result.T)]).T
+    jnp_fun = lambda x: jnp.argwhere(x, size=size, fill_value=fill_value)
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
+    self._CompileAndCheck(jnp_fun, args_maker)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "{}_inshape={}_axis={}".format(
           rec.test_name.capitalize(),