⚡️ Speed up function _check_core_dims by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for _check_core_dims in xarray/core/computation.py

⏱️ Runtime : 3.97 milliseconds → 3.51 milliseconds (best of 5 runs)

📝 Explanation and details

The optimization achieves a 13% speedup by eliminating redundant computations and improving string building efficiency.

Key optimizations:

1. Avoided duplicate set difference calculations: The original code computed set(core_dims) - set(variable_arg.dims) twice - once in the condition and again when building error messages. The optimized version calculates missing_dims once and stores it, reusing the result.

2. Pre-computed missing dimensions: Instead of storing [i, variable_arg, core_dims] and recalculating the set difference later, the optimized version stores the already-computed missing_dims in the tuple, eliminating redundant work during error message construction.

3. Efficient string building: Replaced inefficient string concatenation (message += f"...") with list collection and str.join(), which is significantly faster for building multi-part strings in Python.

4. Reduced list operations: Changed missing += [[...]] (list concatenation) to missing.append((...) (single append), avoiding unnecessary intermediate list creation.

Performance impact by test type:

• Error cases benefit most: Tests with missing dimensions show 5-20% improvements (e.g., test_large_core_dims_some_missing at 19.3% faster) because they avoid the duplicate set calculations during error message generation.
• Success cases see modest gains: Tests where all dimensions are present show smaller but consistent improvements (1-4%) from reduced overhead.
• Large-scale error scenarios: The optimization shines with many missing dimensions (test_large_number_of_variables_all_missing at 42.6% faster) due to the cumulative effect of avoiding repeated calculations.

Context significance: Based on the function reference, _check_core_dims is called within apply_dict_of_variables_vfunc for every variable name in a loop. This makes the optimization particularly valuable since it's in a hot path that processes multiple variables, amplifying the performance gains across xarray operations that involve dimension checking.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

# imports
import pytest
from xarray.core.computation import _check_core_dims


# Helper class to simulate variables with dims property
class DummyVar:
    def __init__(self, dims):
        self.dims = dims

    def __repr__(self):
        return f"DummyVar(dims={self.dims})"


# Helper class to simulate signature object
class DummySignature:
    def __init__(self, input_core_dims):
        self.input_core_dims = input_core_dims


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


def test_all_core_dims_present_single_arg():
    # One variable, all required core dims present
    sig = DummySignature([["x", "y"]])
    var = DummyVar(["x", "y", "z"])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 3.16μs -> 3.40μs (7.12% slower)


def test_all_core_dims_present_multiple_args():
    # Multiple variables, all required core dims present
    sig = DummySignature([["x"], ["y"]])
    var1 = DummyVar(["x", "a"])
    var2 = DummyVar(["y", "b"])
    codeflash_output = _check_core_dims(sig, [var1, var2], "bar")
    result = codeflash_output  # 3.48μs -> 3.42μs (1.61% faster)


def test_no_core_dims_required():
    # No core dims required by signature
    sig = DummySignature([[], []])
    var1 = DummyVar(["a"])
    var2 = DummyVar(["b"])
    codeflash_output = _check_core_dims(sig, [var1, var2], "baz")
    result = codeflash_output  # 3.08μs -> 3.01μs (2.23% faster)


def test_variable_without_dims_property():
    # Variable is a plain object without 'dims' property (should be skipped)
    sig = DummySignature([["x"]])
    var = 123  # int, no dims property
    codeflash_output = _check_core_dims(sig, [var], "nodims")
    result = codeflash_output  # 1.83μs -> 1.86μs (1.93% slower)


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


def test_missing_single_core_dim():
    # One variable, missing one required core dim
    sig = DummySignature([["x", "y"]])
    var = DummyVar(["x"])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 9.09μs -> 8.53μs (6.55% faster)


def test_multiple_missing_core_dims():
    # One variable, missing multiple core dims
    sig = DummySignature([["x", "y", "z"]])
    var = DummyVar(["x"])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 8.25μs -> 7.86μs (5.08% faster)


def test_multiple_args_some_missing():
    # Multiple variables, one missing required dim
    sig = DummySignature([["x"], ["y"]])
    var1 = DummyVar(["x"])
    var2 = DummyVar(["a"])  # missing 'y'
    codeflash_output = _check_core_dims(sig, [var1, var2], "foo")
    result = codeflash_output  # 8.06μs -> 8.01μs (0.662% faster)


def test_multiple_args_multiple_missing():
    # Multiple variables, both missing required dims
    sig = DummySignature([["x"], ["y"]])
    var1 = DummyVar(["a"])
    var2 = DummyVar(["b"])
    codeflash_output = _check_core_dims(sig, [var1, var2], "foo")
    result = codeflash_output  # 10.1μs -> 9.10μs (10.7% faster)


def test_empty_core_dims_and_empty_variable_dims():
    # Variable has empty dims, but no core dims required
    sig = DummySignature([[]])
    var = DummyVar([])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 2.68μs -> 2.89μs (7.10% slower)


def test_variable_with_extra_dims():
    # Variable has more dims than required (should pass)
    sig = DummySignature([["x"]])
    var = DummyVar(["x", "y", "z"])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 2.69μs -> 3.08μs (12.6% slower)


def test_variable_with_no_dims_and_core_dim_required():
    # Variable has dims=[], but signature requires a core dim
    sig = DummySignature([["x"]])
    var = DummyVar([])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 7.77μs -> 7.70μs (0.975% faster)


def test_variable_with_non_string_dims():
    # Variable dims are not strings (should still work)
    sig = DummySignature([[1, 2]])
    var = DummyVar([1, 2, 3])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 2.87μs -> 3.00μs (4.50% slower)


def test_variable_with_partial_overlap_of_dims():
    # Variable has some but not all required dims
    sig = DummySignature([["x", "y"]])
    var = DummyVar(["x", "z"])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 8.41μs -> 8.03μs (4.81% faster)


def test_variable_with_duplicate_dims():
    # Variable has duplicate dims (should not affect set logic)
    sig = DummySignature([["x", "y"]])
    var = DummyVar(["x", "y", "y"])
    codeflash_output = _check_core_dims(sig, [var], "foo")
    result = codeflash_output  # 2.86μs -> 2.99μs (4.38% slower)


def test_signature_longer_than_variable_args():
    # Signature has more entries than variable_args (should only check zipped pairs)
    sig = DummySignature([["x"], ["y"], ["z"]])
    var1 = DummyVar(["x"])
    var2 = DummyVar(["y"])
    codeflash_output = _check_core_dims(sig, [var1, var2], "foo")
    result = codeflash_output  # 3.17μs -> 3.45μs (8.09% slower)


def test_variable_args_longer_than_signature():
    # variable_args has more entries than signature (should only check zipped pairs)
    sig = DummySignature([["x"]])
    var1 = DummyVar(["x"])
    var2 = DummyVar(["y"])
    codeflash_output = _check_core_dims(sig, [var1, var2], "foo")
    result = codeflash_output  # 2.91μs -> 2.81μs (3.41% faster)


def test_variable_with_none_dims():
    # Variable has dims=None (should handle gracefully)
    class DummyVarNoneDims:
        def __init__(self):
            self.dims = None

        def __repr__(self):
            return "DummyVarNoneDims(dims=None)"

    sig = DummySignature([["x"]])
    var = DummyVarNoneDims()
    # set(None) raises TypeError, so should fail
    with pytest.raises(TypeError):
        _check_core_dims(sig, [var], "foo")  # 3.42μs -> 3.22μs (5.99% faster)


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


def test_large_number_of_variables_all_pass():
    # 500 variables, each with a unique required core dim, all present
    n = 500
    sig = DummySignature([[f"dim{i}"] for i in range(n)])
    vars = [DummyVar([f"dim{i}", f"extra{i}"]) for i in range(n)]
    codeflash_output = _check_core_dims(sig, vars, "large_pass")
    result = codeflash_output  # 125μs -> 135μs (7.74% slower)


def test_large_number_of_variables_some_missing():
    # 500 variables, every 10th variable missing its required core dim
    n = 500
    sig = DummySignature([[f"dim{i}"] for i in range(n)])
    vars = []
    for i in range(n):
        if i % 10 == 0:
            # omit the required core dim
            vars.append(DummyVar([f"extra{i}"]))
        else:
            vars.append(DummyVar([f"dim{i}", f"extra{i}"]))
    codeflash_output = _check_core_dims(sig, vars, "large_miss")
    result = codeflash_output  # 188μs -> 179μs (5.37% faster)


def test_large_core_dims_per_variable():
    # Each variable requires 10 core dims, all present
    n = 100
    core_dims = [[f"dim{i}_{j}" for j in range(10)] for i in range(n)]
    sig = DummySignature(core_dims)
    vars = [DummyVar(dims) for dims in core_dims]
    codeflash_output = _check_core_dims(sig, vars, "large_core_dims")
    result = codeflash_output  # 94.9μs -> 91.4μs (3.88% faster)


def test_large_core_dims_some_missing():
    # Each variable requires 10 core dims, but every 5th variable is missing one
    n = 100
    core_dims = [[f"dim{i}_{j}" for j in range(10)] for i in range(n)]
    sig = DummySignature(core_dims)
    vars = []
    for i, dims in enumerate(core_dims):
        if i % 5 == 0:
            # omit the last required core dim
            vars.append(DummyVar(dims[:-1]))
        else:
            vars.append(DummyVar(dims))
    codeflash_output = _check_core_dims(sig, vars, "large_core_dims_miss")
    result = codeflash_output  # 144μs -> 121μs (19.3% faster)


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

from __future__ import annotations

# imports
import pytest
from xarray.core.computation import _check_core_dims


class Signature:
    """Helper class to mimic the signature object with input_core_dims attribute."""

    def __init__(self, input_core_dims):
        self.input_core_dims = input_core_dims


class DummyVar:
    """Helper class to mimic variables with dims attribute."""

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


from xarray.core.computation import _check_core_dims

# unit tests

# 1. Basic Test Cases


def test_all_core_dims_present():
    # Each variable has all required core dims
    signature = Signature([["x", "y"], ["z"]])
    var1 = DummyVar(["x", "y", "a"])
    var2 = DummyVar(["z", "b"])
    codeflash_output = _check_core_dims(signature, [var1, var2], "my_var")
    result = codeflash_output  # 3.52μs -> 3.46μs (1.59% faster)


def test_single_core_dim_missing():
    # One variable missing a required core dim
    signature = Signature([["x"], ["y"]])
    var1 = DummyVar(["x"])
    var2 = DummyVar(["z"])  # Missing "y"
    codeflash_output = _check_core_dims(signature, [var1, var2], "test_var")
    result = codeflash_output  # 9.10μs -> 9.03μs (0.742% faster)


def test_no_core_dims_required():
    # No core dims required (empty lists)
    signature = Signature([[], []])
    var1 = DummyVar(["a", "b"])
    var2 = DummyVar(["c"])
    codeflash_output = _check_core_dims(signature, [var1, var2], "empty_core")
    result = codeflash_output  # 2.96μs -> 3.16μs (6.21% slower)


def test_variable_with_extra_dims():
    # Variable has more dims than required, but all required are present
    signature = Signature([["x"], ["y"]])
    var1 = DummyVar(["x", "foo"])
    var2 = DummyVar(["y", "bar"])
    codeflash_output = _check_core_dims(signature, [var1, var2], "extra_dims")
    result = codeflash_output  # 3.60μs -> 3.52μs (2.04% faster)


def test_variable_is_plain_list():
    # Variable is a plain list (no dims attribute)
    signature = Signature([["x"], ["y"]])
    var1 = [1, 2, 3]  # No dims attribute
    var2 = DummyVar(["y"])
    codeflash_output = _check_core_dims(signature, [var1, var2], "plain_list")
    result = codeflash_output  # 2.98μs -> 2.98μs (0.134% faster)


def test_variable_is_int():
    # Variable is an int (no dims attribute)
    signature = Signature([["x"]])
    var1 = 42
    codeflash_output = _check_core_dims(signature, [var1], "plain_int")
    result = codeflash_output  # 1.98μs -> 2.07μs (4.21% slower)


# 2. Edge Test Cases


def test_empty_signature_and_args():
    # Empty signature and variable_args
    signature = Signature([])
    variable_args = []
    codeflash_output = _check_core_dims(signature, variable_args, "empty")
    result = codeflash_output  # 1.46μs -> 1.57μs (7.43% slower)


def test_signature_longer_than_args():
    # Signature has more entries than variable_args
    signature = Signature([["x"], ["y"], ["z"]])
    var1 = DummyVar(["x"])
    var2 = DummyVar(["y"])
    # Should only check zipped pairs, ignore extra signature entry
    codeflash_output = _check_core_dims(signature, [var1, var2], "short_args")
    result = codeflash_output  # 3.35μs -> 3.38μs (0.652% slower)


def test_args_longer_than_signature():
    # variable_args has more entries than signature
    signature = Signature([["x"]])
    var1 = DummyVar(["x"])
    var2 = DummyVar(["y"])
    # Should only check zipped pairs, ignore extra variable_arg
    codeflash_output = _check_core_dims(signature, [var1, var2], "short_sig")
    result = codeflash_output  # 2.72μs -> 2.81μs (3.03% slower)


def test_variable_with_no_dims_attribute():
    # Variable is an object with no dims attribute
    class NoDims:
        pass

    signature = Signature([["x"]])
    var1 = NoDims()
    codeflash_output = _check_core_dims(signature, [var1], "no_dims")
    result = codeflash_output  # 2.09μs -> 2.27μs (7.86% slower)


def test_variable_with_empty_dims():
    # Variable has dims but is an empty list
    signature = Signature([["x"]])
    var1 = DummyVar([])
    codeflash_output = _check_core_dims(signature, [var1], "empty_dims")
    result = codeflash_output  # 8.83μs -> 8.42μs (4.93% faster)


def test_multiple_missing_core_dims():
    # Multiple variables missing core dims
    signature = Signature([["x", "y"], ["z", "w"]])
    var1 = DummyVar(["x"])  # Missing "y"
    var2 = DummyVar(["z"])  # Missing "w"
    codeflash_output = _check_core_dims(signature, [var1, var2], "multi_missing")
    result = codeflash_output  # 10.5μs -> 10.1μs (3.64% faster)


def test_core_dim_is_none():
    # Core dim is None (should be treated as missing)
    signature = Signature([[None]])
    var1 = DummyVar(["x"])
    codeflash_output = _check_core_dims(signature, [var1], "none_core_dim")
    result = codeflash_output  # 7.46μs -> 7.71μs (3.20% slower)


def test_variable_arg_is_none():
    # variable_arg is None (no dims attribute)
    signature = Signature([["x"]])
    var1 = None
    codeflash_output = _check_core_dims(signature, [var1], "none_var")
    result = codeflash_output  # 1.89μs -> 2.04μs (7.44% slower)


def test_core_dims_are_empty_strings():
    # Core dims are empty strings
    signature = Signature([[""]])
    var1 = DummyVar(["x"])
    codeflash_output = _check_core_dims(signature, [var1], "empty_string_core")
    result = codeflash_output  # 8.71μs -> 8.25μs (5.60% faster)


# 3. Large Scale Test Cases


def test_large_number_of_variables_all_present():
    # 1000 variables, all have required core dims
    n = 1000
    signature = Signature([["dim{}".format(i)] for i in range(n)])
    variable_args = [
        DummyVar(["dim{}".format(i), "other{}".format(i)]) for i in range(n)
    ]
    codeflash_output = _check_core_dims(signature, variable_args, "large_all_present")
    result = codeflash_output  # 243μs -> 259μs (6.00% slower)


def test_large_number_of_variables_some_missing():
    # 1000 variables, every 100th variable missing its core dim
    n = 1000
    signature = Signature([["dim{}".format(i)] for i in range(n)])
    variable_args = []
    for i in range(n):
        if i % 100 == 0:
            variable_args.append(DummyVar(["other{}".format(i)]))  # missing core dim
        else:
            variable_args.append(DummyVar(["dim{}".format(i), "other{}".format(i)]))
    codeflash_output = _check_core_dims(signature, variable_args, "large_some_missing")
    result = codeflash_output  # 267μs -> 261μs (2.11% faster)
    # Should mention all missing dims for every 100th variable
    for i in range(0, n, 100):
        pass


def test_large_number_of_variables_all_missing():
    # 1000 variables, none have required core dims
    n = 1000
    signature = Signature([["dim{}".format(i)] for i in range(n)])
    variable_args = [DummyVar(["other{}".format(i)]) for i in range(n)]
    codeflash_output = _check_core_dims(signature, variable_args, "large_all_missing")
    result = codeflash_output  # 1.44ms -> 1.01ms (42.6% faster)
    # Should mention all missing dims
    for i in range(n):
        pass


def test_large_number_of_variables_no_core_dims():
    # 1000 variables, no core dims required
    n = 1000
    signature = Signature([[] for _ in range(n)])
    variable_args = [DummyVar(["dim{}".format(i)]) for i in range(n)]
    codeflash_output = _check_core_dims(signature, variable_args, "large_no_core")
    result = codeflash_output  # 177μs -> 179μs (1.12% slower)


def test_large_number_of_variables_no_dims_attribute():
    # 1000 variables, none have dims attribute
    n = 1000
    signature = Signature([["dim{}".format(i)] for i in range(n)])
    variable_args = [object() for _ in range(n)]
    codeflash_output = _check_core_dims(signature, variable_args, "large_no_dims_attr")
    result = codeflash_output  # 41.8μs -> 44.6μs (6.23% slower)


# 4. Mutation-sensitive (negative) test: If the function returns True when dims are missing, fail.
def test_mutation_sensitivity_missing_dims_returns_true_should_fail():
    signature = Signature([["x"]])
    var1 = DummyVar(["y"])
    codeflash_output = _check_core_dims(signature, [var1], "mutation_test")
    result = codeflash_output  # 8.86μs -> 8.20μs (8.08% faster)


# 5. Mutation-sensitive (negative) test: If the function returns error message when all dims present, fail.
def test_mutation_sensitivity_all_dims_present_returns_error_should_fail():
    signature = Signature([["x"]])
    var1 = DummyVar(["x"])
    codeflash_output = _check_core_dims(signature, [var1], "mutation_test_2")
    result = codeflash_output  # 2.82μs -> 2.90μs (2.83% slower)


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

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_xarrayteststest_concat_py_xarrayteststest_computation_py_xarrayteststest_formatting_py_xarray__replay_test_0.py::test_xarray_core_computation__check_core_dims	1.08ms	1.06ms	1.80%✅

To edit these changes git checkout codeflash/optimize-_check_core_dims-miyppfhc and push.