Clean up header

Author: evhub

coconut/compiler/templates/header.py_template

@@ -771,10 +771,7 @@ Additionally supports Cartesian products of numpy arrays."""
         if iterables:
             it_modules = [_coconut_get_base_module(it) for it in iterables]
             if _coconut.all(mod in _coconut.numpy_modules for mod in it_modules):
-                if _coconut.any(mod in _coconut.xarray_modules for mod in it_modules):
-                    iterables = tuple((_coconut_xarray_to_numpy(it) if mod in _coconut.xarray_modules else it) for it, mod in _coconut.zip(iterables, it_modules))
-                if _coconut.any(mod in _coconut.pandas_modules for mod in it_modules):
-                    iterables = tuple((it.to_numpy() if mod in _coconut.pandas_modules else it) for it, mod in _coconut.zip(iterables, it_modules))
+                iterables = tuple((it.to_numpy() if mod in _coconut.pandas_modules else _coconut_xarray_to_numpy(it) if mod in _coconut.xarray_modules else it) for it, mod in _coconut.zip(iterables, it_modules))
                 if _coconut.any(mod in _coconut.jax_numpy_modules for mod in it_modules):
                     from jax import numpy
                 else:
@@ -1104,12 +1101,7 @@ class multi_enumerate(_coconut_has_iter):
     through inner iterables and produces a tuple index representing the index
     in each inner iterable. Supports indexing.
 
-    For numpy arrays, effectively equivalent to:
-        it = np.nditer(iterable, flags=["multi_index", "refs_ok"])
-        for x in it:
-            yield it.multi_index, x
-
-    Also supports len for numpy arrays.
+    For numpy arrays, uses np.nditer under the hood and supports len.
     """
     __slots__ = ()
     def __repr__(self):
@@ -1960,10 +1952,10 @@ def all_equal(iterable, to=_coconut_sentinel):
     """
     iterable_module = _coconut_get_base_module(iterable)
     if iterable_module in _coconut.numpy_modules:
-        if iterable_module in _coconut.xarray_modules:
-            iterable = _coconut_xarray_to_numpy(iterable)
-        elif iterable_module in _coconut.pandas_modules:
+        if iterable_module in _coconut.pandas_modules:
             iterable = iterable.to_numpy()
+        elif iterable_module in _coconut.xarray_modules:
+            iterable = _coconut_xarray_to_numpy(iterable)
         return not _coconut.len(iterable) or (iterable == (iterable[0] if to is _coconut_sentinel else to)).all()
     first_item = to
     for item in iterable:

coconut/constants.py

@@ -178,22 +178,22 @@ def get_path_env_var(env_var, default):
     sys.setrecursionlimit(default_recursion_limit)
 
 # modules that numpy-like arrays can live in
-xarray_modules = (
-    "xarray",
+jax_numpy_modules = (
+    "jaxlib",
 )
 pandas_modules = (
     "pandas",
 )
-jax_numpy_modules = (
-    "jaxlib",
+xarray_modules = (
+    "xarray",
 )
 numpy_modules = (
     "numpy",
     "torch",
 ) + (
-    xarray_modules
+    jax_numpy_modules
     + pandas_modules
-    + jax_numpy_modules
+    + xarray_modules
 )
 
 legal_indent_chars = " \t"  # the only Python-legal indent chars