Merge pull request #19245 from jakevdp:complex-warning

PiperOrigin-RevId: 596633005
jax-ml · Jan 8, 2024 · 25cf5cf · 25cf5cf
2 parents 754b74d + 6ab9398
commit 25cf5cf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
@@ -2382,7 +2382,7 @@ def _convert_elt_type_folding_rule(consts, eqn):
       isinstance(o.aval, core.UnshapedArray) and not np.shape(c) and
       not dtypes.issubdtype(eqn.params['new_dtype'], dtypes.extended)):
     with warnings.catch_warnings():
-      warnings.simplefilter('ignore', np.ComplexWarning)
+      warnings.simplefilter('ignore', util.NumpyComplexWarning)
       out = np.array(c).astype(eqn.params['new_dtype'])
     if not o.aval.weak_type:
       return [out], None