⚡️ Speed up function in_stateless_scope by 72%

[codeflash-ai]

📄 72% (0.72x) speedup for in_stateless_scope in keras/src/backend/common/stateless_scope.py

⏱️ Runtime : 828 microseconds → 480 microseconds (best of 169 runs)

📝 Explanation and details

The optimization replaces getattr(GLOBAL_STATE_TRACKER, name, None) with GLOBAL_STATE_TRACKER.__dict__.get(name, None) in the get_global_attribute function, providing a 72% speedup.

Key optimization:

• Direct dictionary lookup instead of Python's reflection mechanism (getattr)
• Bypasses the overhead of Python's attribute resolution protocol and default value handling in C code
• Uses the faster dict.get() method which is optimized at the C level

Why this is faster:
getattr() involves multiple layers of Python's attribute resolution machinery, including descriptor protocol checks and special method lookups. In contrast, threading.local() objects store their per-thread data in a simple __dict__, so direct dictionary access via .get() is much more efficient.

Impact on workloads:
The function references show in_stateless_scope() is called frequently in Keras variable operations - during variable initialization, value access, and assignment operations. Since these are core operations that can occur thousands of times during model training/inference, this micro-optimization has significant cumulative impact.

Test case performance:
The annotated tests show consistent 50-88% speedups across all scenarios, with the optimization being particularly effective for:

• Repeated attribute lookups (74.8% faster in the 1000-iteration test)
• Variable state checking in hot paths
• Both when attributes exist and when they're missing (consistent performance gains)

This optimization is safe because threading.local().__dict__ is the documented way to access thread-local storage and maintains identical behavior while being substantially faster.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import threading

# imports
import pytest  # used for our unit tests
from keras.src.backend.common.stateless_scope import in_stateless_scope

# --- Function to test and its dependencies ---

# Simulate keras/src/backend/common/global_state.py
GLOBAL_STATE_TRACKER = threading.local()

def set_global_attribute(name, value):
    setattr(GLOBAL_STATE_TRACKER, name, value)
from keras.src.backend.common.stateless_scope import in_stateless_scope

# --- Unit tests ---

# Helper to reset stateless_scope between tests
def reset_stateless_scope():
    if hasattr(GLOBAL_STATE_TRACKER, "stateless_scope"):
        delattr(GLOBAL_STATE_TRACKER, "stateless_scope")

# ---------------- Basic Test Cases ----------------

def test_stateless_scope_none_by_default():
    # By default, stateless_scope is not set
    codeflash_output = not in_stateless_scope() # 1.93μs -> 1.54μs (24.9% faster)

def test_stateless_scope_set_true():
    # Setting stateless_scope to True should make in_stateless_scope True
    set_global_attribute("stateless_scope", True)
    codeflash_output = in_stateless_scope() # 1.41μs -> 827ns (70.0% faster)

def test_stateless_scope_set_false():
    # Setting stateless_scope to False should make in_stateless_scope True (since it's not None)
    set_global_attribute("stateless_scope", False)
    codeflash_output = in_stateless_scope() # 1.46μs -> 829ns (76.4% faster)

def test_stateless_scope_set_non_bool():
    # Setting stateless_scope to a non-bool value (e.g., string) should make in_stateless_scope True
    set_global_attribute("stateless_scope", "active")
    codeflash_output = in_stateless_scope() # 1.41μs -> 776ns (81.8% faster)

def test_stateless_scope_set_none_explicitly():
    # Explicitly setting stateless_scope to None should make in_stateless_scope False
    set_global_attribute("stateless_scope", None)
    codeflash_output = not in_stateless_scope() # 1.43μs -> 828ns (72.7% faster)

# ---------------- Edge Test Cases ----------------

def test_stateless_scope_unset_after_set():
    # Set stateless_scope, then unset it, in_stateless_scope should reflect changes
    set_global_attribute("stateless_scope", True)
    codeflash_output = in_stateless_scope() # 1.44μs -> 859ns (67.5% faster)
    reset_stateless_scope()
    codeflash_output = not in_stateless_scope() # 708ns -> 395ns (79.2% faster)

def test_stateless_scope_set_to_empty_string():
    # Setting stateless_scope to empty string should make in_stateless_scope True
    set_global_attribute("stateless_scope", "")
    codeflash_output = in_stateless_scope() # 1.35μs -> 785ns (71.6% faster)

def test_stateless_scope_set_to_zero():
    # Setting stateless_scope to 0 should make in_stateless_scope True (since 0 is not None)
    set_global_attribute("stateless_scope", 0)
    codeflash_output = in_stateless_scope() # 1.34μs -> 795ns (68.9% faster)

def test_stateless_scope_set_to_object():
    # Setting stateless_scope to an object should make in_stateless_scope True
    set_global_attribute("stateless_scope", object())
    codeflash_output = in_stateless_scope() # 1.40μs -> 800ns (75.4% faster)

def test_stateless_scope_set_to_list():
    # Setting stateless_scope to a list should make in_stateless_scope True
    set_global_attribute("stateless_scope", [1, 2, 3])
    codeflash_output = in_stateless_scope() # 1.22μs -> 806ns (51.5% faster)

def test_stateless_scope_set_to_empty_list():
    # Setting stateless_scope to an empty list should make in_stateless_scope True
    set_global_attribute("stateless_scope", [])
    codeflash_output = in_stateless_scope() # 1.37μs -> 757ns (81.4% faster)

def test_stateless_scope_set_to_dict():
    # Setting stateless_scope to a dict should make in_stateless_scope True
    set_global_attribute("stateless_scope", {"active": True})
    codeflash_output = in_stateless_scope() # 1.38μs -> 839ns (64.5% faster)

def test_stateless_scope_set_to_empty_dict():
    # Setting stateless_scope to an empty dict should make in_stateless_scope True
    set_global_attribute("stateless_scope", {})
    codeflash_output = in_stateless_scope() # 1.38μs -> 834ns (64.9% faster)

def test_stateless_scope_set_to_float_nan():
    # Setting stateless_scope to float('nan') should make in_stateless_scope True
    set_global_attribute("stateless_scope", float('nan'))
    codeflash_output = in_stateless_scope() # 1.42μs -> 847ns (67.1% faster)

def test_stateless_scope_set_to_float_inf():
    # Setting stateless_scope to float('inf') should make in_stateless_scope True
    set_global_attribute("stateless_scope", float('inf'))
    codeflash_output = in_stateless_scope() # 1.41μs -> 747ns (88.6% faster)

def test_stateless_scope_set_to_float_minus_inf():
    # Setting stateless_scope to float('-inf') should make in_stateless_scope True
    set_global_attribute("stateless_scope", float('-inf'))
    codeflash_output = in_stateless_scope() # 1.36μs -> 834ns (63.4% faster)

def test_stateless_scope_set_to_bytes():
    # Setting stateless_scope to bytes should make in_stateless_scope True
    set_global_attribute("stateless_scope", b"bytes")
    codeflash_output = in_stateless_scope() # 1.34μs -> 782ns (71.4% faster)

def test_stateless_scope_set_to_tuple():
    # Setting stateless_scope to a tuple should make in_stateless_scope True
    set_global_attribute("stateless_scope", (1, 2, 3))
    codeflash_output = in_stateless_scope() # 1.40μs -> 751ns (87.0% faster)

def test_stateless_scope_set_to_empty_tuple():
    # Setting stateless_scope to an empty tuple should make in_stateless_scope True
    set_global_attribute("stateless_scope", ())
    codeflash_output = in_stateless_scope() # 1.38μs -> 830ns (66.4% faster)

def test_stateless_scope_set_to_none_then_other():
    # Setting stateless_scope to None, then to a value, should update in_stateless_scope
    set_global_attribute("stateless_scope", None)
    codeflash_output = not in_stateless_scope() # 1.41μs -> 838ns (68.0% faster)
    set_global_attribute("stateless_scope", 42)
    codeflash_output = in_stateless_scope() # 673ns -> 390ns (72.6% faster)

def test_stateless_scope_set_to_other_then_none():
    # Setting stateless_scope to a value, then to None, should update in_stateless_scope
    set_global_attribute("stateless_scope", 42)
    codeflash_output = in_stateless_scope() # 1.32μs -> 754ns (74.9% faster)
    set_global_attribute("stateless_scope", None)
    codeflash_output = not in_stateless_scope() # 657ns -> 376ns (74.7% faster)

def test_stateless_scope_many_values():
    # Set stateless_scope to many different values in sequence, always non-None
    for i in range(1000):
        set_global_attribute("stateless_scope", i)
        codeflash_output = in_stateless_scope() # 484μs -> 277μs (74.8% faster)
    # Finally set to None, should be False
    set_global_attribute("stateless_scope", None)
    codeflash_output = not in_stateless_scope() # 497ns -> 282ns (76.2% faster)

def test_stateless_scope_large_data_object():
    # Set stateless_scope to a large list (under 1000 elements)
    large_list = list(range(999))
    set_global_attribute("stateless_scope", large_list)
    codeflash_output = in_stateless_scope() # 2.23μs -> 1.58μs (40.8% faster)
    # Set to None after
    set_global_attribute("stateless_scope", None)
    codeflash_output = not in_stateless_scope() # 635ns -> 363ns (74.9% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import threading

# imports
import pytest  # used for our unit tests
from keras.src.backend.common.stateless_scope import in_stateless_scope

# function to test
# Simulate keras/src/backend/common/global_state.py
GLOBAL_STATE_TRACKER = threading.local()

def set_global_attribute(name, value):
    setattr(GLOBAL_STATE_TRACKER, name, value)
from keras.src.backend.common.stateless_scope import in_stateless_scope

# unit tests

# BASIC TEST CASES

def test_stateless_scope_none_by_default():
    """Test that in_stateless_scope returns False when stateless_scope is not set."""
    # Ensure stateless_scope is unset
    if hasattr(GLOBAL_STATE_TRACKER, "stateless_scope"):
        delattr(GLOBAL_STATE_TRACKER, "stateless_scope")
    codeflash_output = in_stateless_scope() # 1.33μs -> 885ns (50.6% faster)

def test_stateless_scope_set_true():
    """Test that in_stateless_scope returns True when stateless_scope is set to True."""
    set_global_attribute("stateless_scope", True)
    codeflash_output = in_stateless_scope() # 1.44μs -> 817ns (75.8% faster)

def test_stateless_scope_set_false():
    """Test that in_stateless_scope returns True when stateless_scope is set to False (since not None)."""
    set_global_attribute("stateless_scope", False)
    codeflash_output = in_stateless_scope() # 1.43μs -> 812ns (75.7% faster)

def test_stateless_scope_set_non_bool():
    """Test that in_stateless_scope returns True when stateless_scope is set to a non-bool value."""
    set_global_attribute("stateless_scope", "active")
    codeflash_output = in_stateless_scope() # 1.34μs -> 801ns (67.4% faster)

def test_stateless_scope_set_none():
    """Test that in_stateless_scope returns False when stateless_scope is explicitly set to None."""
    set_global_attribute("stateless_scope", None)
    codeflash_output = in_stateless_scope() # 1.39μs -> 824ns (68.4% faster)

# EDGE TEST CASES

def test_stateless_scope_set_empty_string():
    """Test that in_stateless_scope returns True when stateless_scope is set to empty string (not None)."""
    set_global_attribute("stateless_scope", "")
    codeflash_output = in_stateless_scope() # 1.37μs -> 759ns (80.8% faster)

def test_stateless_scope_set_zero():
    """Test that in_stateless_scope returns True when stateless_scope is set to 0 (not None)."""
    set_global_attribute("stateless_scope", 0)
    codeflash_output = in_stateless_scope() # 1.38μs -> 855ns (61.1% faster)

def test_stateless_scope_set_object():
    """Test that in_stateless_scope returns True when stateless_scope is set to an object."""
    class Dummy: pass
    dummy = Dummy()
    set_global_attribute("stateless_scope", dummy)
    codeflash_output = in_stateless_scope() # 1.51μs -> 849ns (78.0% faster)

def test_stateless_scope_set_list():
    """Test that in_stateless_scope returns True when stateless_scope is set to a list."""
    set_global_attribute("stateless_scope", [1,2,3])
    codeflash_output = in_stateless_scope() # 1.42μs -> 808ns (75.4% faster)

def test_stateless_scope_set_dict():
    """Test that in_stateless_scope returns True when stateless_scope is set to a dict."""
    set_global_attribute("stateless_scope", {"key": "value"})
    codeflash_output = in_stateless_scope() # 1.41μs -> 806ns (75.2% faster)

# LARGE SCALE TEST CASES

def test_stateless_scope_massive_set_and_unset():
    """
    Test that repeatedly setting and unsetting stateless_scope works as expected.
    """
    # Unset first
    if hasattr(GLOBAL_STATE_TRACKER, "stateless_scope"):
        delattr(GLOBAL_STATE_TRACKER, "stateless_scope")
    for i in range(100):  # Reasonable upper bound for large scale
        set_global_attribute("stateless_scope", i)
        codeflash_output = in_stateless_scope() # 50.0μs -> 28.9μs (72.6% faster)
        set_global_attribute("stateless_scope", None)
        codeflash_output = in_stateless_scope()

def test_stateless_scope_large_object():
    """
    Test that in_stateless_scope handles large objects as values.
    """
    large_list = [0] * 999
    set_global_attribute("stateless_scope", large_list)
    codeflash_output = in_stateless_scope() # 1.53μs -> 911ns (67.9% faster)
    set_global_attribute("stateless_scope", None)
    codeflash_output = in_stateless_scope() # 612ns -> 342ns (78.9% faster)

def test_stateless_scope_multiple_attributes():
    """
    Test that other attributes do not affect in_stateless_scope.
    """
    set_global_attribute("other_attribute", True)
    if hasattr(GLOBAL_STATE_TRACKER, "stateless_scope"):
        delattr(GLOBAL_STATE_TRACKER, "stateless_scope")
    codeflash_output = in_stateless_scope() # 1.31μs -> 760ns (72.0% faster)
    set_global_attribute("stateless_scope", "active")
    codeflash_output = in_stateless_scope() # 595ns -> 342ns (74.0% faster)

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_integration_testsdataset_testsfashion_mnist_test_py_integration_testsdataset_testscifar10_tes__replay_test_0.py::test_keras_src_backend_common_stateless_scope_in_stateless_scope	66.9μs	40.2μs	66.5%✅
test_pytest_integration_teststorch_workflow_test_py_integration_testsdataset_testscalifornia_housing_test__replay_test_0.py::test_keras_src_backend_common_stateless_scope_in_stateless_scope	124μs	74.2μs	68.0%✅

To edit these changes git checkout codeflash/optimize-in_stateless_scope-mirm1t39 and push.