Raise an error in np.var when array is complex and dtype is not (#2288)

Co-authored-by: vlad <veryfakemail@ya.ru>

jax/numpy/lax_numpy.py

@@ -1646,6 +1646,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):

tests/lax_numpy_test.py

@@ -2665,10 +2665,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(