Repository: pydata/xarray

Issue:
# Better support for symbolic (e.g. jax tracer) values in non-index coordinates

### Is your feature request related to a problem?

Hello! [xarray_jax](https://github.com/google-deepmind/xarray_jax) maintainer here.

I am trying to decide whether it is feasible to rely on using symbolic arrays (such as jax tracers used during JIT compilation) as coordinate data in xarray.

For index coordinates this doesn't work at all, because xarray requires dynamic access to the data in the index to do alignment. So our approach in xarray_jax is to treat index coordinates as static data. I think this is acceptable: even if we could do index alignment with tracers, the result would be variable-shape which jax.jit can't handle without re-tracing anyway.

For non-index coordinates though, it is possible to use jax tracers! However these coordinates have a habit of disappearing due to failed alignment checks, for example in arithmetic operations:

```python
import xarray as xr
import jax
import jax.numpy as jnp

def foo(coord):
   a = xr.DataArray(data=jnp.ones(3), dims=['x'], coords={'foo': (['x'], coord)})
   b = xr.DataArray(data=jnp.ones(3), dims=['x'], coords={'foo': (['x'], coord)})
   print((a + b).coords)

foo(jnp.arange(3))
# prints:
# Coordinates:
#    foo      (x) int32 12B ...

jax.jit(foo)(jnp.arange(3))
# prints:
# Coordinates:
#    *empty*
```

The coordinate is (correctly) preserved when not JIT-ing, because xarray is able to do an equality comparison to check that the coordinate is compatible.

When JITing however, the coordinate array will be a jax Tracer, and the equality comparison will fail with a `jax.errors.TracerBoolConversionError` (or similar). xarray catches this under the hood [here](https://github.com/pydata/xarray/blob/main/xarray/core/variable.py#L1880) since `TracerBoolConversionError` subclasses `TypeError`. xarray then treats the failed comparison as not-equal resulting in the coordinate being dropped.

This isn't entirely consistent with behaviour elsewhere in xarray.

For example when doing an explicit merge it is possible to set `compat='override'` which ensures that coordinates aren't tested for equality and so jax tracers survive. And I believe `compat='override'` is to become the default soon.

`xr.align` also lets jax Tracers survive:

```python
def foo(coord):
   a = xr.DataArray(data=jnp.ones(3), dims=['x'], coords={'foo': (['x'], coord)})
   b = xr.DataArray(data=jnp.ones(3), dims=['x'], coords={'foo': (['x'], coord)})
   a, b = xr.align(a, b)
   print(a.coords)
   print(b.coords)

jax.jit(foo)(jnp.arange(3))
# prints:
# Coordinates:
#     foo      (x) int32 12B ...
# Coordinates:
#     foo      (x) int32 12B ...
```

Overall, while I can see why some might want them, these compatibility checks on non-index coordinates seem a bit of an awkward niche behaviour which is configurable in some places but not others, and (at least after the compat='override' default takes over) will have different defaults in different places too IIUC, which is not ideal. And of course they throw a bit of a spanner in the works of using jax here.

### Describe the solution you'd like

Three suggestions:

1. Add a `xarray.set_options(arithmetic_compat='override')` setting which controls this, and make 'override' the default for it at the same time `compat='override'` becomes the default more widely. Note there is already an `arithmetic_join` setting which controls how index coordinates are joined/aligned, but this doesn't give any control over these compatibility checks on non-index coordinates.
One might also consider exposing the `compat` option on `xr.align` too and any other places where compatibility checks are happening or for consistency ought to happen, to give control over this.

2. Allow `jax.errors.TracerBoolConversionError` to bubble up from generic equality-testing code (or perhaps catch and reraise it as something less jax-specific like `xr.SymbolicComparisonError`) rather than catching it and treating it as unequal. Then in coordinate compatibility-checking code, catch this error and fall back on `compat='override'` behavior in this specific case. This has the advantage that non-index-coord compatibility checks don't have to be disabled across the board, only for jax Tracers. It would require some jax-specific special case in xarray though.
 
3. Just to get rid of alignment checks on non-index coordinates entirely, if noone feels strongly about them. They are already somewhat inconsistent both in configurability and in default behaviour. In many cases it's a safe assumption that any non-index coordinates are a fixed function of the index coordinates and so it's sufficient to check for alignment of the index coordinates. And perhaps other cases it would be OK to make any checking the user's responsibility. I realise this would be backwards-incompatible though.

### Describe alternatives you've considered

See above.

### Additional context

_No response_

Implementation patch:
diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py
index 501c58657cf..85a9d97abeb 100644
--- a/xarray/core/coordinates.py
+++ b/xarray/core/coordinates.py
@@ -21,7 +21,14 @@
     assert_no_index_corrupted,
     create_default_index_implicit,
 )
-from xarray.core.types import DataVars, ErrorOptions, Self, T_DataArray, T_Xarray
+from xarray.core.types import (
+    CompatOptions,
+    DataVars,
+    ErrorOptions,
+    Self,
+    T_DataArray,
+    T_Xarray,
+)
 from xarray.core.utils import (
     Frozen,
     ReprObject,
@@ -31,6 +38,7 @@
 from xarray.core.variable import Variable, as_variable, calculate_dimensions
 from xarray.structure.alignment import Aligner
 from xarray.structure.merge import merge_coordinates_without_align, merge_coords
+from xarray.util.deprecation_helpers import CombineKwargDefault
 
 if TYPE_CHECKING:
     from xarray.core.common import DataWithCoords
@@ -500,18 +508,20 @@ def _drop_coords(self, coord_names):
         # redirect to DatasetCoordinates._drop_coords
         self._data.coords._drop_coords(coord_names)
 
-    def _merge_raw(self, other, reflexive):
+    def _merge_raw(self, other, reflexive, compat: CompatOptions | CombineKwargDefault):
         """For use with binary arithmetic."""
         if other is None:
             variables = dict(self.variables)
             indexes = dict(self.xindexes)
         else:
             coord_list = [self, other] if not reflexive else [other, self]
-            variables, indexes = merge_coordinates_without_align(coord_list)
+            variables, indexes = merge_coordinates_without_align(
+                coord_list, compat=compat
+            )
         return variables, indexes
 
     @contextmanager
-    def _merge_inplace(self, other):
+    def _merge_inplace(self, other, compat: CompatOptions | CombineKwargDefault):
         """For use with in-place binary arithmetic."""
         if other is None:
             yield
@@ -524,12 +534,17 @@ def _merge_inplace(self, other):
                 if k not in self.xindexes
             }
             variables, indexes = merge_coordinates_without_align(
-                [self, other], prioritized
+                [self, other], prioritized, compat=compat
             )
             yield
             self._update_coords(variables, indexes)
 
-    def merge(self, other: Mapping[Any, Any] | None) -> Dataset:
+    def merge(
+        self,
+        other: Mapping[Any, Any] | None,
+        *,
+        compat: CompatOptions | CombineKwargDefault = "minimal",
+    ) -> Dataset:
         """Merge two sets of coordinates to create a new Dataset
 
         The method implements the logic used for joining coordinates in the
@@ -546,6 +561,8 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset:
         other : dict-like, optional
             A :py:class:`Coordinates` object or any mapping that can be turned
             into coordinates.
+        compat : {"identical", "equals", "broadcast_equals", "no_conflicts", "override", "minimal"}, default: "minimal"
+            Compatibility checks to use between coordinate variables.
 
         Returns
         -------
@@ -560,7 +577,7 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset:
         if not isinstance(other, Coordinates):
             other = Dataset(coords=other).coords
 
-        coords, indexes = merge_coordinates_without_align([self, other])
+        coords, indexes = merge_coordinates_without_align([self, other], compat=compat)
         coord_names = set(coords)
         return Dataset._construct_direct(
             variables=coords, coord_names=coord_names, indexes=indexes
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
index fcfa0317131..706c35e5459 100644
--- a/xarray/core/dataarray.py
+++ b/xarray/core/dataarray.py
@@ -4906,7 +4906,9 @@ def _binary_op(
             if not reflexive
             else f(other_variable_or_arraylike, self.variable)
         )
-        coords, indexes = self.coords._merge_raw(other_coords, reflexive)
+        coords, indexes = self.coords._merge_raw(
+            other_coords, reflexive, compat=OPTIONS["arithmetic_compat"]
+        )
         name = result_name([self, other])
 
         return self._replace(variable, coords, name, indexes=indexes)
@@ -4926,7 +4928,9 @@ def _inplace_binary_op(self, other: DaCompatible, f: Callable) -> Self:
         other_coords = getattr(other, "coords", None)
         other_variable = getattr(other, "variable", other)
         try:
-            with self.coords._merge_inplace(other_coords):
+            with self.coords._merge_inplace(
+                other_coords, compat=OPTIONS["arithmetic_compat"]
+            ):
                 f(self.variable, other_variable)
         except MergeError as exc:
             raise MergeError(
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index 84a67d95412..e15f1077639 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -7763,7 +7763,7 @@ def apply_over_both(lhs_data_vars, rhs_data_vars, lhs_vars, rhs_vars):
             return type(self)(new_data_vars)
 
         other_coords: Coordinates | None = getattr(other, "coords", None)
-        ds = self.coords.merge(other_coords)
+        ds = self.coords.merge(other_coords, compat=OPTIONS["arithmetic_compat"])
 
         if isinstance(other, Dataset):
             new_vars = apply_over_both(
diff --git a/xarray/core/options.py b/xarray/core/options.py
index 451070ce7b4..9af6117d237 100644
--- a/xarray/core/options.py
+++ b/xarray/core/options.py
@@ -2,14 +2,16 @@
 
 import warnings
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Literal, TypedDict
+from typing import TYPE_CHECKING, Any, Literal, TypedDict, get_args
 
+from xarray.core.types import CompatOptions
 from xarray.core.utils import FrozenDict
 
 if TYPE_CHECKING:
     from matplotlib.colors import Colormap
 
     Options = Literal[
+        "arithmetic_compat",
         "arithmetic_join",
         "chunk_manager",
         "cmap_divergent",
@@ -40,6 +42,7 @@
 
     class T_Options(TypedDict):
         arithmetic_broadcast: bool
+        arithmetic_compat: CompatOptions
         arithmetic_join: Literal["inner", "outer", "left", "right", "exact"]
         chunk_manager: str
         cmap_divergent: str | Colormap
@@ -70,6 +73,7 @@ class T_Options(TypedDict):
 
 OPTIONS: T_Options = {
     "arithmetic_broadcast": True,
+    "arithmetic_compat": "minimal",
     "arithmetic_join": "inner",
     "chunk_manager": "dask",
     "cmap_divergent": "RdBu_r",
@@ -109,6 +113,7 @@ def _positive_integer(value: Any) -> bool:
 
 _VALIDATORS = {
     "arithmetic_broadcast": lambda value: isinstance(value, bool),
+    "arithmetic_compat": get_args(CompatOptions).__contains__,
     "arithmetic_join": _JOIN_OPTIONS.__contains__,
     "display_max_children": _positive_integer,
     "display_max_rows": _positive_integer,
@@ -178,8 +183,27 @@ class set_options:
 
     Parameters
     ----------
+    arithmetic_broadcast : bool, default: True
+        Whether to perform automatic broadcasting in binary operations.
+    arithmetic_compat: {"identical", "equals", "broadcast_equals", "no_conflicts", "override", "minimal"}, default: "minimal"
+        How to compare non-index coordinates of the same name for potential
+        conflicts when performing binary operations. (For the alignment of index
+        coordinates in binary operations, see `arithmetic_join`.)
+
+        - "identical": all values, dimensions and attributes of the coordinates
+          must be the same.
+        - "equals": all values and dimensions of the coordinates must be the
+          same.
+        - "broadcast_equals": all values of the coordinates must be equal after
+          broadcasting to ensure common dimensions.
+        - "no_conflicts": only values which are not null in both coordinates
+          must be equal. The returned coordinate then contains the combination
+          of all non-null values.
+        - "override": skip comparing and take the coordinates from the first
+          operand.
+        - "minimal": drop conflicting coordinates.
     arithmetic_join : {"inner", "outer", "left", "right", "exact"}, default: "inner"
-        DataArray/Dataset alignment in binary operations:
+        DataArray/Dataset index alignment in binary operations:
 
         - "outer": use the union of object indexes
         - "inner": use the intersection of object indexes
@@ -187,9 +211,6 @@ class set_options:
         - "right": use indexes from the last object with each dimension
         - "exact": instead of aligning, raise `ValueError` when indexes to be
           aligned are not equal
-        - "override": if indexes are of same size, rewrite indexes to be
-          those of the first object with that dimension. Indexes for the same
-          dimension must have the same size in all objects.
     chunk_manager : str, default: "dask"
         Chunk manager to use for chunked array computations when multiple
         options are installed.
diff --git a/xarray/structure/merge.py b/xarray/structure/merge.py
index d398a37c8ef..d4790421070 100644
--- a/xarray/structure/merge.py
+++ b/xarray/structure/merge.py
@@ -438,6 +438,7 @@ def merge_coordinates_without_align(
     prioritized: Mapping[Any, MergeElement] | None = None,
     exclude_dims: AbstractSet = frozenset(),
     combine_attrs: CombineAttrsOptions = "override",
+    compat: CompatOptions | CombineKwargDefault = "minimal",
 ) -> tuple[dict[Hashable, Variable], dict[Hashable, Index]]:
     """Merge variables/indexes from coordinates without automatic alignments.
 
@@ -462,7 +463,7 @@ def merge_coordinates_without_align(
     # TODO: indexes should probably be filtered in collected elements
     # before merging them
     merged_coords, merged_indexes = merge_collected(
-        filtered, prioritized, combine_attrs=combine_attrs
+        filtered, prioritized, compat=compat, combine_attrs=combine_attrs
     )
     merged_indexes = filter_indexes_from_coords(merged_indexes, set(merged_coords))
 


Test patch:
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index d1aacba6aaa..71fd7ec6cee 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -25,6 +25,7 @@
     DataArray,
     Dataset,
     IndexVariable,
+    MergeError,
     Variable,
     align,
     broadcast,
@@ -2626,6 +2627,43 @@ def test_math_with_coords(self) -> None:
         actual = alt + orig
         assert_identical(expected, actual)
 
+    def test_math_with_arithmetic_compat_options(self) -> None:
+        # Setting up a clash of non-index coordinate 'foo':
+        a = xr.DataArray(
+            data=[0, 0, 0],
+            dims=["x"],
+            coords={
+                "x": [1, 2, 3],
+                "foo": (["x"], [1.0, 2.0, np.nan]),
+            },
+        )
+        b = xr.DataArray(
+            data=[0, 0, 0],
+            dims=["x"],
+            coords={
+                "x": [1, 2, 3],
+                "foo": (["x"], [np.nan, 2.0, 3.0]),
+            },
+        )
+
+        with xr.set_options(arithmetic_compat="minimal"):
+            assert_equal(a + b, a.drop_vars("foo"))
+
+        with xr.set_options(arithmetic_compat="override"):
+            assert_equal(a + b, a)
+            assert_equal(b + a, b)
+
+        with xr.set_options(arithmetic_compat="no_conflicts"):
+            expected = a.assign_coords(foo=(["x"], [1.0, 2.0, 3.0]))
+            assert_equal(a + b, expected)
+            assert_equal(b + a, expected)
+
+        with xr.set_options(arithmetic_compat="equals"):
+            with pytest.raises(MergeError):
+                a + b
+            with pytest.raises(MergeError):
+                b + a
+
     def test_index_math(self) -> None:
         orig = DataArray(range(3), dims="x", name="x")
         actual = orig + 1
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 4835c562d5a..a11bf55b28f 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -6990,6 +6990,41 @@ def test_binary_op_join_setting(self) -> None:
             actual = ds1 + ds2
             assert_equal(actual, expected)
 
+    def test_binary_op_compat_setting(self) -> None:
+        # Setting up a clash of non-index coordinate 'foo':
+        a = xr.Dataset(
+            data_vars={"var": (["x"], [0, 0, 0])},
+            coords={
+                "x": [1, 2, 3],
+                "foo": (["x"], [1.0, 2.0, np.nan]),
+            },
+        )
+        b = xr.Dataset(
+            data_vars={"var": (["x"], [0, 0, 0])},
+            coords={
+                "x": [1, 2, 3],
+                "foo": (["x"], [np.nan, 2.0, 3.0]),
+            },
+        )
+
+        with xr.set_options(arithmetic_compat="minimal"):
+            assert_equal(a + b, a.drop_vars("foo"))
+
+        with xr.set_options(arithmetic_compat="override"):
+            assert_equal(a + b, a)
+            assert_equal(b + a, b)
+
+        with xr.set_options(arithmetic_compat="no_conflicts"):
+            expected = a.assign_coords(foo=(["x"], [1.0, 2.0, 3.0]))
+            assert_equal(a + b, expected)
+            assert_equal(b + a, expected)
+
+        with xr.set_options(arithmetic_compat="equals"):
+            with pytest.raises(MergeError):
+                a + b
+            with pytest.raises(MergeError):
+                b + a
+
     @pytest.mark.parametrize(
         ["keep_attrs", "expected"],
         (
diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py
index 0cd888f5782..cb9e81da97a 100644
--- a/xarray/tests/test_datatree.py
+++ b/xarray/tests/test_datatree.py
@@ -2423,6 +2423,46 @@ def test_arithmetic_inherited_coords(self) -> None:
         expected["/foo/bar"].data = np.array([8, 10, 12])
         assert_identical(actual, expected)
 
+    def test_binary_op_compat_setting(self) -> None:
+        # Setting up a clash of non-index coordinate 'foo':
+        a = DataTree(
+            xr.Dataset(
+                data_vars={"var": (["x"], [0, 0, 0])},
+                coords={
+                    "x": [1, 2, 3],
+                    "foo": (["x"], [1.0, 2.0, np.nan]),
+                },
+            )
+        )
+        b = DataTree(
+            xr.Dataset(
+                data_vars={"var": (["x"], [0, 0, 0])},
+                coords={
+                    "x": [1, 2, 3],
+                    "foo": (["x"], [np.nan, 2.0, 3.0]),
+                },
+            )
+        )
+
+        with xr.set_options(arithmetic_compat="minimal"):
+            expected = DataTree(a.dataset.drop_vars("foo"))
+            assert_equal(a + b, expected)
+
+        with xr.set_options(arithmetic_compat="override"):
+            assert_equal(a + b, a)
+            assert_equal(b + a, b)
+
+        with xr.set_options(arithmetic_compat="no_conflicts"):
+            expected = DataTree(a.dataset.assign_coords(foo=(["x"], [1.0, 2.0, 3.0])))
+            assert_equal(a + b, expected)
+            assert_equal(b + a, expected)
+
+        with xr.set_options(arithmetic_compat="equals"):
+            with pytest.raises(xr.MergeError):
+                a + b
+            with pytest.raises(xr.MergeError):
+                b + a
+
     def test_binary_op_commutativity_with_dataset(self) -> None:
         # regression test for #9365