xarray
xarray copied to clipboard
Fix assert_equal with check_dim_order=False for mixed dimension orders
Summary
- Fixes #10704
assert_equalwithcheck_dim_order=Falsenow works correctly for Datasets containing variables with different dimension orders
The Bug
When comparing Datasets with check_dim_order=False, the comparison would fail if individual variables had different dimension orders, even when comparing a Dataset to itself:
dataset = xr.Dataset({
"foo": xr.DataArray(np.zeros([4, 5]), dims=("a", "b")),
"bar": xr.DataArray(np.zeros([5, 4]), dims=("b", "a"))
})
# This would fail, but shouldn't:
xr.testing.assert_equal(dataset, dataset, check_dim_order=False)
The Fix
The solution transposes both objects to a canonical dimension order using the intersection of their dimensions. The ellipsis (...) handles any dimensions unique to either object:
common_dims = set(a.dims) & set(b.dims)
if common_dims:
canonical_order = sorted(common_dims) + [...]
return a.transpose(*canonical_order), b.transpose(*canonical_order)
This elegant approach works uniformly for Variables, DataArrays, and Datasets without special casing.
Test plan
- [x] Added test case
test_assert_equal_dataset_check_dim_orderthat reproduces the issue - [x] All existing tests pass
- [x] Verified fix works for both
assert_equalandassert_allclose
🤖 Generated with Claude Code