Repository: google/jax

Issue:
np.var does not match onp.var in some cases
Running this snippet
```python
import jax.numpy as np 
import numpy as onp 
inp_array = np.array([[[5.29215704-0.30965656j,  1.20047163+1.23179551j, 0+0.43213071j,  6.7226796 +4.36282052j]],   
                      [[5.60267397+2.28311318j,           0+0.36502505j,         0 +1j, -0.45407162+1.00102298j]]])  
print(np.var(inp_array, 2, dtype=np.float32)) 
print(onp.var(onp.array(inp_array), 2, dtype=onp.float32))
```
Prints different outputs for numpy and jax.numpy
```
[[10.911064]
 [ 6.728332]]
[[12.953884]
 [ 8.079251]]
```
Not sure if there is an actual use case for np.var on complex arrays with dtype=float, but I guess the outputs are supposed to be the same since jax.numpy is supposed to match numpy where it's possible.

The mismatch happens because the implementation in NumPy uses dtype for all the calculations, meanwhile, jax.numpy only converts the computed value to dtype in the very end and does all the calculations in the promote_types(complex, dtype=float) which evaluates to complex.

If this sounds like something that should be fixed then I can make a PR.


Discussion:
I think we should do one of two things:
(a) match NumPy, or
(b) issue an error, if the behavior isn't sensible.

I can't particularly decide if the behavior of NumPy is sensible here; it might make more sense to issue an error.

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
@@ -1399,6 +1399,14 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
 
   a_dtype = _dtype(a)
   if dtype:
+    if (not issubdtype(dtype, complexfloating) and
+        issubdtype(a_dtype, complexfloating)):
+      msg = ("jax.numpy.var does not yet support real dtype parameters when "
+             "computing the variance of an array of complex values. The "
+             "semantics of numpy.var seem unclear in this case. Please comment "
+             "on https://github.com/google/jax/issues/2283 if this behavior is "
+             "important to you.")
+      raise ValueError(msg)
     a_dtype = promote_types(a_dtype, dtype)
   else:
     if not issubdtype(a_dtype, inexact):


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
@@ -2307,10 +2307,14 @@ def onp_fun(x):
     jnp_fun = partial(jnp.var, dtype=out_dtype, axis=axis, ddof=ddof, keepdims=keepdims)
     tol = jtu.tolerance(out_dtype, {onp.float16: 1e-1, onp.float32: 1e-3,
                                     onp.float64: 1e-3, onp.complex128: 1e-6})
-    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True,
-                            tol=tol)
-    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, rtol=tol,
-                          atol=tol)
+    if (jnp.issubdtype(dtype, jnp.complexfloating) and
+        not jnp.issubdtype(out_dtype, jnp.complexfloating)):
+      self.assertRaises(ValueError, lambda: jnp_fun(*args_maker()))
+    else:
+      self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True,
+                              tol=tol)
+      self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, rtol=tol,
+                            atol=tol)
 
   @parameterized.named_parameters(
       jtu.cases_from_list(