Repository: google/jax

Issue:
Improper kwarg forwarding for np.mean
Steps to reproduce

```
import jax.numpy as np
import numpy as onp
x = lax.add(np.eye(3), 0.)
onp.mean(x)
```

Traceback:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-91-45d207ceb919> in <module>()
      1 x = lax.add(np.eye(3), 0.)
----> 2 onp.mean(x)

~/anaconda/lib/python3.5/site-packages/numpy/core/fromnumeric.py in mean(a, axis, dtype, out, keepdims)
   2915             pass
   2916         else:
-> 2917             return mean(axis=axis, dtype=dtype, out=out, **kwargs)
   2918 
   2919     return _methods._mean(a, axis=axis, dtype=dtype,

TypeError: mean() got an unexpected keyword argument 'out'
```


Discussion:
Discussing with @dougalm, this might be intended behavior or a bug on NumPy's part? 
Calling `onp.mean(x.copy())` works, so it might be NumPy not properly using DeviceArray's `__array__` method.

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
@@ -596,7 +596,7 @@ def _make_reduction(np_fun, op, init_val, preproc=None):
   @_wraps(np_fun)
   def reduction(a, axis=None, dtype=None, out=None, keepdims=False):
     if out is not None:
-      raise ValueError("reduction does not support `out` argument.")
+      raise ValueError("reduction does not support the `out` argument.")
 
     a = a if isinstance(a, ndarray) else asarray(a)
     a = preproc(a) if preproc else a
@@ -644,7 +644,10 @@ def _reduction_init_val(a, init_val):
 
 
 @_wraps(onp.mean)
-def mean(a, axis=None, dtype=None, keepdims=False):
+def mean(a, axis=None, dtype=None, out=None, keepdims=False):
+  if out is not None:
+    raise ValueError("mean does not support the `out` argument.")
+
   if axis is None:
     normalizer = size(a)
   else:
@@ -663,7 +666,10 @@ def mean(a, axis=None, dtype=None, keepdims=False):
 
 
 @_wraps(onp.var)
-def var(a, axis=None, dtype=None, keepdims=False, ddof=0):
+def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
+  if out is not None:
+    raise ValueError("var does not support the `out` argument.")
+
   if ddof != 0:
     raise NotImplementedError("Only implemented for ddof=0.")
   if dtype is None:
@@ -677,8 +683,10 @@ def var(a, axis=None, dtype=None, keepdims=False, ddof=0):
 
 
 @_wraps(onp.std)
-def std(a, axis=None, dtype=None, keepdims=False, ddof=0):
-  return sqrt(var(a, axis=axis, dtype=dtype, keepdims=keepdims, ddof=ddof))
+def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
+  if out is not None:
+    raise ValueError("std does not support the `out` argument.")
+  return sqrt(var(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims))
 
 
 @_wraps(onp.allclose)


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
@@ -27,6 +27,7 @@
 import numpy as onp
 
 from jax import api
+from jax import lax
 from jax import numpy as lnp
 from jax import test_util as jtu
 
@@ -810,6 +811,12 @@ def testRavel(self):
 
   # TODO(mattjj): test other ndarray-like method overrides
 
+  def testOnpMean(self):
+    # from https://github.com/google/jax/issues/125
+    x = lax.add(lnp.eye(3), 0.)
+    ans = onp.mean(x)
+    self.assertAllClose(ans, onp.array([1./3, 1./3, 1./3]), check_dtypes=False)
+
 
 if __name__ == "__main__":
   absltest.main()