⚡️ Speed up function is_layer_block by 67%

[codeflash-ai]

📄 67% (0.67x) speedup for is_layer_block in src/transformers/model_debugging_utils.py

⏱️ Runtime : 532 microseconds → 318 microseconds (best of 105 runs)

📝 Explanation and details

The optimization replaces the any() function with an explicit for loop that returns True immediately upon finding the first match. This provides a 67% speedup through early termination and reduced Python function call overhead.

Key changes:

• Replaced any(f".{number}." in child.get("module_path", "") for child in node["children"]) with a manual loop that breaks early
• Pre-computed the search string f".{number}." outside the loop to avoid repeated string formatting
• Used explicit return True on first match instead of evaluating all children

Why this is faster:

1. Early termination: The original any() creates a generator expression that Python must fully set up even if the first element matches. The manual loop exits immediately on the first match.
2. Reduced overhead: Eliminates the function call overhead of any() and the generator setup cost.
3. String formatting optimization: Computing search_str = f".{number}." once instead of recreating it in each iteration.

Performance impact based on function_references:
The is_layer_block function is called within prune_intermediate_layers(), which processes model debugging trees. Since this likely runs on large model architectures with many layer blocks, the optimization becomes significant when processing hundreds of children nodes.

Test case analysis:

• Best improvements (71-102% faster): Large-scale tests where no children match, as early termination prevents scanning all remaining children
• Moderate improvements (30-50% faster): Cases with mixed or matching children benefit from reduced overhead
• Minimal improvements (3-20% faster): Edge cases with few children see smaller but consistent gains

The optimization is particularly effective for transformer models with deep layer hierarchies where most layer blocks don't match the pattern.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import re

# imports
from transformers.model_debugging_utils import is_layer_block


# function to test
# Copyright 2025 The HuggingFace Inc. team.
# All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


LAYER_SUFFIX_RE = re.compile(r"(.*)\.(\d+)$")  # should be generic enough, ends with a number

# unit tests

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


def test_basic_true_single_child():
    # Basic: module_path ends with number, has one child with correct module_path
    node = {
        "module_path": "model.encoder.layer.2",
        "children": [
            {"module_path": "model.encoder.layer.2.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.51μs -> 2.49μs (40.9% faster)


def test_basic_false_no_children():
    # Basic: module_path ends with number, but no children
    node = {"module_path": "model.encoder.layer.2", "children": []}
    codeflash_output = is_layer_block(node)  # 1.66μs -> 1.38μs (20.4% faster)


def test_basic_false_no_number_suffix():
    # Basic: module_path does not end with a number
    node = {
        "module_path": "model.encoder.layer",
        "children": [
            {"module_path": "model.encoder.layer.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 1.56μs -> 1.34μs (16.6% faster)


def test_basic_true_multiple_children():
    # Basic: multiple children, one matches the .{number}. pattern
    node = {
        "module_path": "model.encoder.layer.5",
        "children": [
            {"module_path": "model.encoder.layer.5.attention"},
            {"module_path": "model.encoder.layer.5.feedforward"},
            {"module_path": "model.encoder.layer.4.attention"},  # should not match
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.46μs -> 2.44μs (41.8% faster)


def test_basic_false_children_do_not_match():
    # Basic: children exist, but none match the .{number}. pattern
    node = {
        "module_path": "model.encoder.layer.3",
        "children": [
            {"module_path": "model.encoder.layer.2.attention"},
            {"module_path": "model.encoder.layer.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.06μs -> 2.35μs (30.4% faster)


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


def test_edge_module_path_empty_string():
    # Edge: module_path is empty string
    node = {
        "module_path": "",
        "children": [
            {"module_path": "model.encoder.layer.0.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 1.03μs -> 857ns (20.2% faster)


def test_edge_children_none():
    # Edge: children is None
    node = {"module_path": "model.encoder.layer.1", "children": None}
    codeflash_output = is_layer_block(node)  # 2.09μs -> 2.02μs (3.62% faster)


def test_edge_children_missing_key():
    # Edge: children key is missing
    node = {"module_path": "model.encoder.layer.1"}
    codeflash_output = is_layer_block(node)  # 1.71μs -> 1.51μs (13.1% faster)


def test_edge_child_module_path_missing():
    # Edge: child dict does not have module_path key
    node = {
        "module_path": "model.encoder.layer.1",
        "children": [{"not_module_path": "model.encoder.layer.1.attention"}],
    }
    codeflash_output = is_layer_block(node)  # 3.66μs -> 2.80μs (30.6% faster)


def test_edge_module_path_number_at_start():
    # Edge: module_path starts with a number but not at the end
    node = {
        "module_path": "2.model.encoder.layer",
        "children": [
            {"module_path": "2.model.encoder.layer.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 1.67μs -> 1.52μs (9.93% faster)


def test_edge_module_path_multiple_numbers():
    # Edge: module_path ends with a number, but has numbers elsewhere too
    node = {
        "module_path": "model.10.encoder.layer.2",
        "children": [
            {"module_path": "model.10.encoder.layer.2.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.60μs -> 2.48μs (45.4% faster)


def test_edge_child_module_path_number_in_middle():
    # Edge: child module_path has number in middle, not matching pattern
    node = {
        "module_path": "model.encoder.layer.2",
        "children": [
            {"module_path": "model.encoder.2.layer.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.27μs -> 2.17μs (50.6% faster)


def test_edge_module_path_with_trailing_dot():
    # Edge: module_path ends with dot and number
    node = {
        "module_path": "model.encoder.layer.3.",
        "children": [
            {"module_path": "model.encoder.layer.3.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 1.83μs -> 1.68μs (9.19% faster)


def test_edge_child_module_path_exact_match():
    # Edge: child module_path is exactly "model.encoder.layer.2"
    node = {
        "module_path": "model.encoder.layer.2",
        "children": [
            {"module_path": "model.encoder.layer.2"},
        ],
    }
    # Should be False, since ".2." is not present in child module_path
    codeflash_output = is_layer_block(node)  # 3.30μs -> 2.35μs (40.7% faster)


def test_edge_child_module_path_partial_match():
    # Edge: child module_path contains ".2." but not at expected position
    node = {
        "module_path": "model.encoder.layer.2",
        "children": [
            {"module_path": "model.encoder.layer.2extra.attention"},
        ],
    }
    # Should be False, ".2." is not present as substring
    codeflash_output = is_layer_block(node)  # 3.13μs -> 2.30μs (36.1% faster)


def test_edge_module_path_number_is_zero():
    # Edge: module_path ends with 0
    node = {
        "module_path": "model.encoder.layer.0",
        "children": [
            {"module_path": "model.encoder.layer.0.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.32μs -> 2.36μs (40.8% faster)


def test_edge_module_path_number_is_large():
    # Edge: module_path ends with a large number
    node = {
        "module_path": "model.encoder.layer.999",
        "children": [
            {"module_path": "model.encoder.layer.999.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.35μs -> 2.40μs (39.7% faster)


def test_edge_child_module_path_is_empty_string():
    # Edge: child module_path is empty string
    node = {
        "module_path": "model.encoder.layer.1",
        "children": [
            {"module_path": ""},
        ],
    }
    codeflash_output = is_layer_block(node)  # 2.92μs -> 2.06μs (41.8% faster)


def test_large_scale_many_children_one_match():
    # Large scale: many children, only one matches the pattern
    children = [{"module_path": f"model.encoder.layer.1.child{i}"} for i in range(999)]
    children.append({"module_path": "model.encoder.layer.1.attention"})
    node = {"module_path": "model.encoder.layer.1", "children": children}
    codeflash_output = is_layer_block(node)  # 4.34μs -> 3.30μs (31.6% faster)


def test_large_scale_all_children_match():
    # Large scale: all children match the pattern
    children = [{"module_path": f"model.encoder.layer.5.child{i}.sub"} for i in range(1000)]
    node = {"module_path": "model.encoder.layer.5", "children": children}
    # All children have ".5." in their module_path
    codeflash_output = is_layer_block(node)  # 3.91μs -> 2.86μs (36.7% faster)


def test_large_scale_no_children_match():
    # Large scale: none of the children match the pattern
    children = [{"module_path": f"model.encoder.layer.4.child{i}.sub"} for i in range(1000)]
    node = {"module_path": "model.encoder.layer.5", "children": children}
    codeflash_output = is_layer_block(node)  # 69.5μs -> 40.6μs (71.2% faster)


def test_large_scale_empty_children():
    # Large scale: children is empty list
    node = {"module_path": "model.encoder.layer.5", "children": []}
    codeflash_output = is_layer_block(node)  # 1.84μs -> 1.60μs (15.3% faster)


def test_large_scale_module_path_large_number():
    # Large scale: module_path ends with a large number, children match
    children = [{"module_path": f"model.encoder.layer.999.child{i}.sub"} for i in range(1000)]
    node = {"module_path": "model.encoder.layer.999", "children": children}
    codeflash_output = is_layer_block(node)  # 4.05μs -> 2.79μs (44.9% faster)


def test_large_scale_module_path_large_number_no_match():
    # Large scale: module_path ends with a large number, children do not match
    children = [{"module_path": f"model.encoder.layer.998.child{i}.sub"} for i in range(1000)]
    node = {"module_path": "model.encoder.layer.999", "children": children}
    codeflash_output = is_layer_block(node)  # 70.7μs -> 40.8μs (73.4% faster)


def test_large_scale_children_with_missing_module_path():
    # Large scale: some children do not have module_path key
    children = [{"module_path": f"model.encoder.layer.5.child{i}.sub"} for i in range(995)]
    children += [{"not_module_path": f"model.encoder.layer.5.child{i}.sub"} for i in range(5)]
    node = {"module_path": "model.encoder.layer.5", "children": children}
    codeflash_output = is_layer_block(node)  # 3.67μs -> 2.57μs (42.8% faster)


def test_large_scale_children_all_missing_module_path():
    # Large scale: all children missing module_path key
    children = [{"not_module_path": f"model.encoder.layer.5.child{i}.sub"} for i in range(1000)]
    node = {"module_path": "model.encoder.layer.5", "children": children}
    codeflash_output = is_layer_block(node)  # 62.3μs -> 30.8μs (102% faster)


def test_large_scale_module_path_with_multiple_numbers():
    # Large scale: module_path contains multiple numbers, ends with one
    children = [{"module_path": f"model.10.encoder.layer.7.child{i}.sub"} for i in range(1000)]
    node = {"module_path": "model.10.encoder.layer.7", "children": children}
    codeflash_output = is_layer_block(node)  # 3.85μs -> 2.67μs (43.9% faster)


def test_large_scale_children_are_empty_dicts():
    # Large scale: children are empty dicts
    children = [{} for _ in range(1000)]
    node = {"module_path": "model.encoder.layer.7", "children": children}
    codeflash_output = is_layer_block(node)  # 60.8μs -> 30.1μs (102% faster)


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

import re

# imports
from transformers.model_debugging_utils import is_layer_block


# function to test
# Copyright 2025 The HuggingFace Inc. team.
# All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


LAYER_SUFFIX_RE = re.compile(r"(.*)\.(\d+)$")  # should be generic enough, ends with a number

# unit tests

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


def test_basic_layer_block_true():
    # Node with module_path ending in .number, children with .number. in module_path
    node = {
        "module_path": "model.encoder.layer.5",
        "children": [
            {"module_path": "model.encoder.layer.5.attention"},
            {"module_path": "model.encoder.layer.5.output"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.70μs -> 2.68μs (38.1% faster)


def test_basic_layer_block_false_no_children():
    # Node with module_path ending in .number, but no children
    node = {"module_path": "model.encoder.layer.5", "children": []}
    codeflash_output = is_layer_block(node)  # 1.66μs -> 1.54μs (7.68% faster)


def test_basic_layer_block_false_no_number_suffix():
    # Node with module_path not ending in .number
    node = {
        "module_path": "model.encoder.layer.attention",
        "children": [
            {"module_path": "model.encoder.layer.attention.output"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 1.56μs -> 1.57μs (0.828% slower)


def test_basic_layer_block_false_children_wrong_number():
    # Node with module_path ending in .number, but children with wrong number in module_path
    node = {
        "module_path": "model.encoder.layer.3",
        "children": [
            {"module_path": "model.encoder.layer.4.attention"},
            {"module_path": "model.encoder.layer.2.output"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.30μs -> 2.39μs (37.9% faster)


def test_basic_layer_block_true_mixed_children():
    # Node with module_path ending in .number, children with both correct and incorrect numbers
    node = {
        "module_path": "model.encoder.layer.1",
        "children": [
            {"module_path": "model.encoder.layer.1.attention"},
            {"module_path": "model.encoder.layer.2.output"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.21μs -> 2.17μs (47.4% faster)


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


def test_edge_empty_module_path():
    # Node with empty module_path
    node = {
        "module_path": "",
        "children": [
            {"module_path": "model.encoder.layer.0.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 992ns -> 896ns (10.7% faster)


def test_edge_missing_module_path_key():
    # Node missing 'module_path' key
    node = {
        "children": [
            {"module_path": "model.encoder.layer.0.attention"},
        ]
    }
    codeflash_output = is_layer_block(node)  # 1.39μs -> 1.24μs (12.6% faster)


def test_edge_missing_children_key():
    # Node missing 'children' key
    node = {"module_path": "model.encoder.layer.0"}
    codeflash_output = is_layer_block(node)  # 1.86μs -> 1.62μs (14.4% faster)


def test_edge_children_none():
    # Node with children=None
    node = {"module_path": "model.encoder.layer.0", "children": None}
    codeflash_output = is_layer_block(node)  # 1.67μs -> 1.56μs (6.93% faster)


def test_edge_child_missing_module_path():
    # Child missing 'module_path' key
    node = {
        "module_path": "model.encoder.layer.0",
        "children": [
            {},
            {"module_path": "model.encoder.layer.0.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 4.34μs -> 3.16μs (37.6% faster)


def test_edge_module_path_with_multiple_numbers():
    # Node with multiple numbers in module_path, only last should be used
    node = {
        "module_path": "model.2.encoder.layer.3",
        "children": [
            {"module_path": "model.2.encoder.layer.3.attention"},
            {"module_path": "model.2.encoder.layer.2.output"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 4.12μs -> 2.96μs (39.3% faster)


def test_edge_module_path_number_at_start():
    # Node with number at the start, not at the end
    node = {
        "module_path": "1.model.encoder.layer",
        "children": [
            {"module_path": "1.model.encoder.layer.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 1.63μs -> 1.53μs (6.19% faster)


def test_edge_module_path_number_with_leading_zeros():
    # Node with module_path ending in .0001
    node = {
        "module_path": "model.encoder.layer.0001",
        "children": [
            {"module_path": "model.encoder.layer.0001.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.67μs -> 2.59μs (41.7% faster)


def test_edge_module_path_number_negative():
    # Node with module_path ending in .-1 (should not match, as regex only matches digits)
    node = {
        "module_path": "model.encoder.layer.-1",
        "children": [
            {"module_path": "model.encoder.layer.-1.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 1.57μs -> 1.46μs (7.89% faster)


def test_edge_module_path_number_float():
    # Node with module_path ending in .1.5 (should not match)
    node = {
        "module_path": "model.encoder.layer.1.5",
        "children": [
            {"module_path": "model.encoder.layer.1.5.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.41μs -> 2.34μs (45.6% faster)


def test_edge_child_module_path_partial_match():
    # Child module_path contains .number but not as a separate segment
    node = {
        "module_path": "model.encoder.layer.7",
        "children": [
            {"module_path": "model.encoder.layer.77.attention"},
            {"module_path": "model.encoder.layer.7.attention"},
        ],
    }
    codeflash_output = is_layer_block(node)  # 3.33μs -> 2.36μs (41.1% faster)


def test_edge_child_module_path_number_in_middle():
    # Child module_path contains .number. in the middle
    node = {
        "module_path": "model.encoder.layer.8",
        "children": [
            {"module_path": "model.encoder.8.layer.attention"},
        ],
    }
    # Should not match, as the number must be after the last dot in parent
    codeflash_output = is_layer_block(node)  # 3.21μs -> 2.07μs (54.9% faster)


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


def test_large_scale_all_true():
    # Node with module_path ending in .number, 500 children all matching .number.
    number = "42"
    node = {
        "module_path": f"model.encoder.layer.{number}",
        "children": [{"module_path": f"model.encoder.layer.{number}.submodule{i}"} for i in range(500)],
    }
    codeflash_output = is_layer_block(node)  # 3.53μs -> 2.60μs (35.9% faster)


def test_large_scale_mixed_children():
    # Node with module_path ending in .number, 500 children, half matching, half not
    number = "123"
    children = []
    for i in range(250):
        children.append({"module_path": f"model.encoder.layer.{number}.submodule{i}"})
    for i in range(250):
        children.append({"module_path": f"model.encoder.layer.{int(number) + 1}.submodule{i}"})
    node = {"module_path": f"model.encoder.layer.{number}", "children": children}
    codeflash_output = is_layer_block(node)  # 3.44μs -> 2.66μs (29.0% faster)


def test_large_scale_none_matching():
    # Node with module_path ending in .number, 500 children, none matching .number.
    number = "99"
    node = {
        "module_path": f"model.encoder.layer.{number}",
        "children": [{"module_path": f"model.encoder.layer.{int(number) + 1}.submodule{i}"} for i in range(500)],
    }
    codeflash_output = is_layer_block(node)  # 36.5μs -> 21.3μs (71.3% faster)


def test_large_scale_no_children():
    # Node with module_path ending in .number, children is empty list
    node = {"module_path": "model.encoder.layer.0", "children": []}
    codeflash_output = is_layer_block(node)  # 1.70μs -> 1.55μs (9.56% faster)


def test_large_scale_children_missing_module_path():
    # Node with many children, some missing module_path
    number = "5"
    children = [{"module_path": f"model.encoder.layer.{number}.submodule{i}"} for i in range(490)]
    children += [{} for _ in range(10)]
    node = {"module_path": f"model.encoder.layer.{number}", "children": children}
    codeflash_output = is_layer_block(node)  # 3.44μs -> 2.25μs (52.7% faster)


def test_large_scale_children_all_missing_module_path():
    # Node with many children, all missing module_path
    node = {"module_path": "model.encoder.layer.10", "children": [{} for _ in range(500)]}
    codeflash_output = is_layer_block(node)  # 32.0μs -> 16.4μs (95.7% faster)


def test_large_scale_module_path_high_number():
    # Node with module_path ending in a large number, children matching
    number = "999"
    node = {
        "module_path": f"model.encoder.layer.{number}",
        "children": [{"module_path": f"model.encoder.layer.{number}.submodule{i}"} for i in range(999)],
    }
    codeflash_output = is_layer_block(node)  # 3.81μs -> 2.90μs (31.2% faster)


def test_large_scale_module_path_high_number_none_matching():
    # Node with module_path ending in a large number, children not matching
    number = "999"
    node = {
        "module_path": f"model.encoder.layer.{number}",
        "children": [{"module_path": f"model.encoder.layer.{int(number) - 1}.submodule{i}"} for i in range(999)],
    }
    codeflash_output = is_layer_block(node)  # 69.8μs -> 40.0μs (74.4% 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-is_layer_block-misp25r2 and push.