⚡️ Speed up function unified_dim_sizes by 5%

[codeflash-ai]

📄 5% (0.05x) speedup for unified_dim_sizes in xarray/core/computation.py

⏱️ Runtime : 773 microseconds → 735 microseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 5% speedup through several key micro-optimizations targeting the hot path where unified_dim_sizes is called:

Key optimizations:

1. Pre-normalized exclude_dims lookup: Instead of checking membership in the original exclude_dims parameter (which could be any Set type), the code pre-converts it to a set/frozenset for O(1) membership testing. This avoids repeated type checking overhead in the inner loop.

2. Eliminated tuple creation overhead: Replaced zip(var.dims, var.shape) with direct indexing (for i in range(len(dims))), avoiding the creation of temporary tuples for each dimension-size pair.

3. Conditional duplicate detection: Only converts dims to a set when necessary (when len(dims) > 1), avoiding unnecessary set creation for single-dimension variables.

4. Single dictionary lookup with dict.get(): Uses dim_sizes.get(dim) instead of checking dim not in dim_sizes followed by assignment, reducing dictionary lookups from two to one per dimension.

Performance characteristics from tests:

• Shows significant gains (24-91% faster) for scenarios with many variables with disjoint dimensions or no dimensions
• Performs slightly slower (1-25%) on small cases with few dimensions due to optimization overhead
• The function is called from apply_variable_ufunc, a core xarray computation function that processes Variable broadcasting, making these micro-optimizations valuable for data processing workloads

Impact on workloads: Since unified_dim_sizes is used in xarray's core computation pipeline for dimension broadcasting, these optimizations will benefit any operation involving multiple xarray Variables, especially when processing large numbers of variables or performing repeated computations.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 70 Passed
🌀 Generated Regression Tests	✅ 45 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_computation.py::test_unified_dim_sizes	13.7μs	14.0μs	-2.01%⚠️

🌀 Generated Regression Tests and Runtime

from collections.abc import Hashable, Iterable, Set

# imports
import pytest  # used for our unit tests
from xarray.core.computation import unified_dim_sizes


# Minimal Variable class for testing
class Variable:
    """
    Minimal implementation of xarray.core.variable.Variable for unit testing.
    """

    def __init__(self, dims, shape):
        self.dims = tuple(dims)
        self.shape = tuple(shape)


from xarray.core.computation import unified_dim_sizes

# unit tests

# 1. Basic Test Cases


def test_single_variable_single_dim():
    # One variable, one dimension
    var = Variable(dims=["x"], shape=[5])
    codeflash_output = unified_dim_sizes([var])
    result = codeflash_output  # 2.90μs -> 2.94μs (1.33% slower)


def test_single_variable_multi_dims():
    # One variable, multiple dimensions
    var = Variable(dims=["x", "y"], shape=[5, 3])
    codeflash_output = unified_dim_sizes([var])
    result = codeflash_output  # 2.92μs -> 3.29μs (11.1% slower)


def test_multiple_variables_disjoint_dims():
    # Multiple variables, disjoint dimensions
    var1 = Variable(dims=["x"], shape=[5])
    var2 = Variable(dims=["y"], shape=[3])
    codeflash_output = unified_dim_sizes([var1, var2])
    result = codeflash_output  # 3.25μs -> 2.98μs (9.14% faster)


def test_multiple_variables_shared_dims_same_size():
    # Multiple variables, shared dimension, same size
    var1 = Variable(dims=["x"], shape=[5])
    var2 = Variable(dims=["x", "y"], shape=[5, 3])
    codeflash_output = unified_dim_sizes([var1, var2])
    result = codeflash_output  # 3.40μs -> 3.62μs (6.18% slower)


def test_exclude_dims_basic():
    # Exclude a dimension
    var1 = Variable(dims=["x"], shape=[5])
    var2 = Variable(dims=["x", "y"], shape=[5, 3])
    codeflash_output = unified_dim_sizes([var1, var2], exclude_dims={"x"})
    result = codeflash_output  # 3.54μs -> 3.60μs (1.58% slower)


# 2. Edge Test Cases


def test_duplicate_dims_in_variable_raises():
    # Variable with duplicate dimension names
    var = Variable(dims=["x", "x"], shape=[5, 5])
    with pytest.raises(ValueError, match="duplicate"):
        unified_dim_sizes([var])  # 4.03μs -> 4.35μs (7.38% slower)


def test_mismatched_dim_sizes_raises():
    # Variables with same dim but different sizes
    var1 = Variable(dims=["x"], shape=[5])
    var2 = Variable(dims=["x"], shape=[3])
    with pytest.raises(ValueError, match="mismatched lengths"):
        unified_dim_sizes([var1, var2])  # 4.63μs -> 4.18μs (10.7% faster)


def test_empty_variables_list():
    # No variables: should return empty dict
    codeflash_output = unified_dim_sizes([])
    result = codeflash_output  # 609ns -> 1.11μs (45.1% slower)


def test_variable_with_no_dims():
    # Variable with no dimensions
    var = Variable(dims=[], shape=[])
    codeflash_output = unified_dim_sizes([var])
    result = codeflash_output  # 2.17μs -> 2.04μs (6.36% faster)


def test_exclude_all_dims():
    # All dims excluded, should return empty dict
    var1 = Variable(dims=["x", "y"], shape=[5, 3])
    codeflash_output = unified_dim_sizes([var1], exclude_dims={"x", "y"})
    result = codeflash_output  # 2.79μs -> 2.96μs (5.72% slower)


def test_non_string_dims():
    # Non-string dimension names (e.g., integers)
    var1 = Variable(dims=[1, 2], shape=[10, 20])
    var2 = Variable(dims=[2], shape=[20])
    codeflash_output = unified_dim_sizes([var1, var2])
    result = codeflash_output  # 3.58μs -> 3.92μs (8.89% slower)


def test_hashable_dims_types():
    # Hashable dimension names (tuples)
    var1 = Variable(dims=[("x", 1), ("y", 2)], shape=[4, 5])
    var2 = Variable(dims=[("y", 2)], shape=[5])
    codeflash_output = unified_dim_sizes([var1, var2])
    result = codeflash_output  # 3.80μs -> 3.81μs (0.262% slower)


def test_exclude_dims_partial_overlap():
    # Exclude dims partially present in variables
    var1 = Variable(dims=["x", "y"], shape=[5, 3])
    var2 = Variable(dims=["y", "z"], shape=[3, 7])
    codeflash_output = unified_dim_sizes([var1, var2], exclude_dims={"y"})
    result = codeflash_output  # 3.43μs -> 3.97μs (13.5% slower)


def test_exclude_dims_not_present():
    # Exclude dims not present in any variable, should have no effect
    var1 = Variable(dims=["x"], shape=[5])
    codeflash_output = unified_dim_sizes([var1], exclude_dims={"not_present"})
    result = codeflash_output  # 2.72μs -> 2.35μs (16.0% faster)


def test_variable_with_zero_shape():
    # Variable with shape zero (empty dimension)
    var = Variable(dims=["x"], shape=[0])
    codeflash_output = unified_dim_sizes([var])
    result = codeflash_output  # 2.49μs -> 2.47μs (0.728% faster)


# 3. Large Scale Test Cases


def test_large_number_of_variables_and_dims():
    # Many variables, many dimensions, all sizes match
    num_vars = 100
    num_dims = 10
    dims = [f"dim{i}" for i in range(num_dims)]
    shape = [i + 1 for i in range(num_dims)]
    variables = [Variable(dims=dims, shape=shape) for _ in range(num_vars)]
    codeflash_output = unified_dim_sizes(variables)
    result = codeflash_output  # 76.0μs -> 102μs (25.6% slower)
    expected = {f"dim{i}": i + 1 for i in range(num_dims)}


def test_large_number_of_variables_with_disjoint_dims():
    # Many variables, each with a unique dimension
    num_vars = 100
    variables = [Variable(dims=[f"dim{i}"], shape=[i + 1]) for i in range(num_vars)]
    codeflash_output = unified_dim_sizes(variables)
    result = codeflash_output  # 34.6μs -> 27.8μs (24.2% faster)
    expected = {f"dim{i}": i + 1 for i in range(num_vars)}


def test_large_number_of_variables_with_mismatched_sizes():
    # Many variables, one dimension, two sizes, should raise
    num_vars = 100
    variables = [Variable(dims=["x"], shape=[5]) for _ in range(num_vars // 2)] + [
        Variable(dims=["x"], shape=[7]) for _ in range(num_vars // 2)
    ]
    with pytest.raises(ValueError, match="mismatched lengths"):
        unified_dim_sizes(variables)  # 16.3μs -> 13.2μs (23.6% faster)


def test_large_exclude_dims():
    # Many variables and large exclude_dims set
    num_vars = 50
    num_dims = 20
    dims = [f"dim{i}" for i in range(num_dims)]
    shape = [i + 2 for i in range(num_dims)]
    variables = [Variable(dims=dims, shape=shape) for _ in range(num_vars)]
    exclude = set(dims[:10])  # exclude half the dims
    codeflash_output = unified_dim_sizes(variables, exclude_dims=exclude)
    result = codeflash_output  # 60.5μs -> 76.5μs (20.9% slower)
    expected = {f"dim{i}": i + 2 for i in range(10, num_dims)}


def test_large_variables_with_no_dims():
    # Many variables, all with no dims
    variables = [Variable(dims=[], shape=[]) for _ in range(500)]
    codeflash_output = unified_dim_sizes(variables)
    result = codeflash_output  # 96.3μs -> 50.4μs (91.0% faster)


def test_large_variables_some_with_no_dims():
    # Mix of variables with and without dims
    variables = [Variable(dims=[], shape=[]) for _ in range(200)] + [
        Variable(dims=["x"], shape=[8]) for _ in range(300)
    ]
    codeflash_output = unified_dim_sizes(variables)
    result = codeflash_output  # 112μs -> 77.0μs (45.7% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from collections.abc import Hashable, Iterable, Set

# imports
import pytest  # used for our unit tests
from xarray.core.computation import unified_dim_sizes


# Minimal Variable class for testing purposes
class Variable:
    """
    A minimal stand-in for xarray.core.variable.Variable.
    """

    def __init__(self, dims, shape):
        self.dims = dims
        self.shape = shape


from xarray.core.computation import unified_dim_sizes

# unit tests

# ------------------------
# BASIC TEST CASES
# ------------------------


def test_single_variable_single_dim():
    # Single variable, one dimension
    v = Variable(("x",), (5,))
    codeflash_output = unified_dim_sizes([v])
    result = codeflash_output  # 2.78μs -> 2.96μs (6.02% slower)


def test_single_variable_multiple_dims():
    # Single variable, multiple dimensions
    v = Variable(("x", "y"), (3, 4))
    codeflash_output = unified_dim_sizes([v])
    result = codeflash_output  # 2.91μs -> 3.21μs (9.22% slower)


def test_multiple_variables_disjoint_dims():
    # Two variables, disjoint dimensions
    v1 = Variable(("x",), (2,))
    v2 = Variable(("y",), (7,))
    codeflash_output = unified_dim_sizes([v1, v2])
    result = codeflash_output  # 3.21μs -> 3.05μs (5.24% faster)


def test_multiple_variables_shared_dim_same_size():
    # Two variables, shared dimension with same size
    v1 = Variable(("x",), (8,))
    v2 = Variable(("x", "y"), (8, 2))
    codeflash_output = unified_dim_sizes([v1, v2])
    result = codeflash_output  # 3.41μs -> 3.77μs (9.64% slower)


def test_multiple_variables_shared_dim_different_sizes_raises():
    # Two variables, shared dimension with different sizes
    v1 = Variable(("x",), (8,))
    v2 = Variable(("x", "y"), (9, 2))
    with pytest.raises(ValueError, match="mismatched lengths for dimension x: 8 vs 9"):
        unified_dim_sizes([v1, v2])  # 4.10μs -> 4.27μs (4.07% slower)


def test_exclude_dims_removes_from_result():
    # Exclude a dimension, should not appear in result
    v1 = Variable(("x",), (5,))
    v2 = Variable(("x", "y"), (5, 3))
    codeflash_output = unified_dim_sizes([v1, v2], exclude_dims={"x"})
    result = codeflash_output  # 3.42μs -> 3.72μs (8.07% slower)


def test_exclude_dims_all_dims():
    # Exclude all dims, should return empty dict
    v1 = Variable(("x", "y"), (2, 3))
    codeflash_output = unified_dim_sizes([v1], exclude_dims={"x", "y"})
    result = codeflash_output  # 2.55μs -> 2.66μs (4.02% slower)


def test_empty_variables_list():
    # No variables, should return empty dict
    codeflash_output = unified_dim_sizes([])
    result = codeflash_output  # 563ns -> 1.09μs (48.3% slower)


# ------------------------
# EDGE TEST CASES
# ------------------------


def test_variable_with_no_dims():
    # Variable with no dims (scalar)
    v = Variable((), ())
    codeflash_output = unified_dim_sizes([v])
    result = codeflash_output  # 2.13μs -> 1.93μs (10.4% faster)


def test_variable_with_duplicate_dims_raises():
    # Variable with duplicate dims should raise
    v = Variable(("x", "x"), (2, 2))
    with pytest.raises(ValueError, match="duplicate.*dimensions.*variable"):
        unified_dim_sizes([v])  # 4.18μs -> 4.37μs (4.32% slower)


def test_shared_dim_with_exclude_dim():
    # Shared dim, but excluded, so no error
    v1 = Variable(("x",), (5,))
    v2 = Variable(("x",), (7,))
    codeflash_output = unified_dim_sizes([v1, v2], exclude_dims={"x"})
    result = codeflash_output  # 3.17μs -> 2.69μs (18.0% faster)


def test_variable_with_non_hashable_dim_raises():
    # Non-hashable dim should raise TypeError
    v = Variable((["x"],), (2,))
    with pytest.raises(TypeError):
        unified_dim_sizes([v])  # 2.05μs -> 3.06μs (32.9% slower)


def test_variable_with_empty_dim_name():
    # Empty string as dim name
    v = Variable(("",), (3,))
    codeflash_output = unified_dim_sizes([v])
    result = codeflash_output  # 2.57μs -> 2.47μs (3.92% faster)


def test_variables_with_mixed_dim_types():
    # Mix string and int as dim names
    v1 = Variable(("x",), (4,))
    v2 = Variable((1,), (5,))
    v3 = Variable(("x", 1), (4, 5))
    codeflash_output = unified_dim_sizes([v1, v2, v3])
    result = codeflash_output  # 4.03μs -> 4.43μs (8.90% slower)


def test_variable_with_zero_length_dim():
    # Variable with zero-length dimension
    v = Variable(("x",), (0,))
    codeflash_output = unified_dim_sizes([v])
    result = codeflash_output  # 2.31μs -> 2.24μs (2.86% faster)


def test_variable_with_large_dim_name():
    # Very long string as dim name
    long_dim = "x" * 100
    v = Variable((long_dim,), (2,))
    codeflash_output = unified_dim_sizes([v])
    result = codeflash_output  # 2.42μs -> 2.36μs (2.41% faster)


def test_variable_with_none_dim_name():
    # None as a dim name is allowed (hashable)
    v = Variable((None,), (3,))
    codeflash_output = unified_dim_sizes([v])
    result = codeflash_output  # 2.66μs -> 2.45μs (8.56% faster)


# ------------------------
# LARGE SCALE TEST CASES
# ------------------------


def test_many_variables_many_dims():
    # 100 variables, each with a unique dimension
    variables = [Variable((f"dim{i}",), (i + 1,)) for i in range(100)]
    codeflash_output = unified_dim_sizes(variables)
    result = codeflash_output  # 34.7μs -> 27.5μs (26.3% faster)
    expected = {f"dim{i}": i + 1 for i in range(100)}


def test_many_variables_shared_dims_same_size():
    # 100 variables, all with the same dimension and size
    variables = [Variable(("x",), (10,)) for _ in range(100)]
    codeflash_output = unified_dim_sizes(variables)
    result = codeflash_output  # 27.4μs -> 20.8μs (31.9% faster)


def test_many_variables_shared_dims_different_sizes_raises():
    # 100 variables, first 50 with size 10, next 50 with size 11
    variables = [Variable(("x",), (10,)) for _ in range(50)] + [
        Variable(("x",), (11,)) for _ in range(50)
    ]
    with pytest.raises(
        ValueError, match="mismatched lengths for dimension x: 10 vs 11"
    ):
        unified_dim_sizes(variables)  # 16.8μs -> 13.0μs (29.3% faster)


def test_large_number_of_dims_per_variable():
    # One variable with 100 dims
    dims = tuple(f"d{i}" for i in range(100))
    shape = tuple(i + 1 for i in range(100))
    v = Variable(dims, shape)
    codeflash_output = unified_dim_sizes([v])
    result = codeflash_output  # 16.2μs -> 18.8μs (13.8% slower)
    expected = {f"d{i}": i + 1 for i in range(100)}


def test_large_mix_exclude_dims():
    # 50 variables, each with 2 dims, exclude half the dims
    variables = [
        Variable((f"dim{i}", f"dim{i+1}"), (i + 1, i + 2)) for i in range(0, 100, 2)
    ]
    exclude = {f"dim{i}" for i in range(0, 100, 4)}
    codeflash_output = unified_dim_sizes(variables, exclude_dims=exclude)
    result = codeflash_output  # 24.5μs -> 27.2μs (9.90% slower)
    # Only dims not in exclude should appear
    expected = {}
    for i in range(0, 100, 2):
        if f"dim{i}" not in exclude:
            expected[f"dim{i}"] = i + 1
        if f"dim{i+1}" not in exclude:
            expected[f"dim{i+1}"] = i + 2


def test_performance_many_variables_many_dims():
    # 100 variables, each with 10 unique dims
    variables = []
    expected = {}
    for i in range(100):
        dims = tuple(f"dim{i}_{j}" for j in range(10))
        shape = tuple(j + 1 for j in range(10))
        variables.append(Variable(dims, shape))
        for j in range(10):
            expected[f"dim{i}_{j}"] = j + 1
    codeflash_output = unified_dim_sizes(variables)
    result = codeflash_output  # 120μs -> 150μs (19.9% slower)


def test_exclude_dims_large_scale_all_excluded():
    # 100 variables, each with 1 dim, all dims excluded
    variables = [Variable((f"dim{i}",), (i + 1,)) for i in range(100)]
    exclude = {f"dim{i}" for i in range(100)}
    codeflash_output = unified_dim_sizes(variables, exclude_dims=exclude)
    result = codeflash_output  # 28.6μs -> 18.5μs (54.7% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-unified_dim_sizes-miyrj69b and push.