⚡️ Speed up function serialize_keras_class_and_config by 55%

[codeflash-ai]

📄 55% (0.55x) speedup for serialize_keras_class_and_config in keras/src/legacy/saving/serialization.py

⏱️ Runtime : 2.68 milliseconds → 1.72 milliseconds (best of 167 runs)

📝 Explanation and details

The optimization achieves a 55% speedup by eliminating redundant function calls to _shared_object_saving_scope().

Key optimizations:

1. Reduced function calls: The original code calls _shared_object_saving_scope() twice when both conditions are met - once in the compound if statement and again when calling scope.get_config(). The optimized version calls it at most once by storing the result in a local variable scope.

2. Short-circuit evaluation: By checking obj is not None first, the optimization avoids calling _shared_object_saving_scope() entirely when no object is provided (which happens in many test cases), saving expensive getattr() operations on the threading.local object.

Performance impact from profiling:

• _shared_object_saving_scope() calls dropped from 4,242 to 1,715 (59% reduction)
• The expensive compound condition line went from 76.9% of total time to being eliminated
• Overall function time reduced from 16.17ms to 10.40ms

Test case benefits:

• Basic cases without objects see 200-260% speedup (most common usage pattern)
• Cases with objects but no scope see modest improvements or slight slowdowns due to the extra obj is not None check
• Large-scale tests with many objects show consistent 130-170% improvements

Real-world impact:
Based on the function reference, serialize_keras_class_and_config is called from serialize_keras_object, which is likely in the hot path during model serialization. The optimization is particularly effective when serializing models without shared object scopes (the common case), where it eliminates unnecessary threading.local lookups entirely.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import threading

# imports
import pytest
from keras.src.legacy.saving.serialization import \
    serialize_keras_class_and_config

# function to test
# (copied from keras/src/legacy/saving/serialization.py)
SHARED_OBJECT_KEY = "shared_object_id"
SHARED_OBJECT_SAVING = threading.local()
from keras.src.legacy.saving.serialization import \
    serialize_keras_class_and_config

# Helper classes for mocking shared object saving scope
class DummyScope:
    """Dummy scope for testing shared object saving behavior."""
    def __init__(self):
        self.configs = {}
        self.created = []
    def get_config(self, obj):
        # Return config if already serialized
        return self.configs.get(id(obj), None)
    def create_config(self, base_config, obj):
        # Save and return config for object
        self.created.append((id(obj), base_config))
        self.configs[id(obj)] = {"created": True, **base_config}
        return self.configs[id(obj)]

# ----------- UNIT TESTS ------------

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

def test_basic_minimal_config():
    # Basic: minimal config with only class name and config
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}); result = codeflash_output # 1.73μs -> 550ns (214% faster)

def test_basic_with_shared_object_id():
    # Basic: config with shared_object_id
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, shared_object_id="abc123"); result = codeflash_output # 2.00μs -> 867ns (131% faster)

def test_basic_with_obj_but_no_scope():
    # Basic: obj is provided, but no scope is set
    obj = object()
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj); result = codeflash_output # 1.82μs -> 2.25μs (18.9% slower)

def test_basic_with_obj_and_shared_object_id():
    # Basic: obj and shared_object_id provided, but no scope
    obj = object()
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj, shared_object_id="id42"); result = codeflash_output # 2.22μs -> 2.31μs (3.73% slower)

def test_basic_config_types():
    # Basic: config can be any mapping, even empty or nested
    codeflash_output = serialize_keras_class_and_config("Sequential", {}); result = codeflash_output # 1.71μs -> 560ns (206% faster)
    codeflash_output = serialize_keras_class_and_config("Model", {"layers": [{"type": "Dense", "units": 10}]}); result2 = codeflash_output # 751ns -> 277ns (171% faster)

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

def test_edge_none_class_name():
    # Edge: class_name is None
    codeflash_output = serialize_keras_class_and_config(None, {"units": 32}); result = codeflash_output # 1.61μs -> 520ns (210% faster)

def test_edge_none_config():
    # Edge: config is None
    codeflash_output = serialize_keras_class_and_config("Dense", None); result = codeflash_output # 1.60μs -> 517ns (210% faster)

def test_edge_empty_strings():
    # Edge: empty string values
    codeflash_output = serialize_keras_class_and_config("", {}); result = codeflash_output # 1.62μs -> 452ns (260% faster)

def test_edge_shared_object_id_is_zero():
    # Edge: shared_object_id is 0 (falsy but valid)
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, shared_object_id=0); result = codeflash_output # 2.10μs -> 916ns (129% faster)

def test_edge_shared_object_id_is_false():
    # Edge: shared_object_id is False (falsy but valid)
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, shared_object_id=False); result = codeflash_output # 2.06μs -> 878ns (135% faster)

def test_edge_config_is_list():
    # Edge: config is a list (not typical, but accepted by function)
    codeflash_output = serialize_keras_class_and_config("Dense", [1, 2, 3]); result = codeflash_output # 1.63μs -> 488ns (234% faster)

def test_edge_config_is_string():
    # Edge: config is a string (not typical, but accepted by function)
    codeflash_output = serialize_keras_class_and_config("Dense", "units=32"); result = codeflash_output # 1.67μs -> 488ns (241% faster)

def test_edge_config_is_int():
    # Edge: config is an int (not typical, but accepted by function)
    codeflash_output = serialize_keras_class_and_config("Dense", 42); result = codeflash_output # 1.68μs -> 495ns (239% faster)

def test_edge_with_scope_get_config_returns_config():
    # Edge: scope returns a config for the object, should return that config
    obj = object()
    scope = DummyScope()
    config_for_obj = {"class_name": "Dense", "config": {"units": 32}, "custom": True}
    scope.configs[id(obj)] = config_for_obj
    SHARED_OBJECT_SAVING.scope = scope
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj); result = codeflash_output # 1.84μs -> 2.23μs (17.3% slower)
    SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_edge_with_scope_get_config_returns_none():
    # Edge: scope returns None, should call create_config
    obj = object()
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj); result = codeflash_output # 1.78μs -> 1.87μs (4.81% slower)
    SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_edge_with_scope_and_shared_object_id():
    # Edge: scope, obj, and shared_object_id all provided
    obj = object()
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj, shared_object_id="id99"); result = codeflash_output # 1.99μs -> 2.06μs (3.39% slower)
    SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_edge_scope_is_set_but_obj_is_none():
    # Edge: scope is set, but obj is None, should ignore scope
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=None); result = codeflash_output # 1.75μs -> 709ns (147% faster)
    SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_edge_scope_is_set_but_obj_is_unhashable():
    # Edge: obj is unhashable, but id(obj) is always available
    class Unhashable:
        __hash__ = None
    obj = Unhashable()
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj); result = codeflash_output # 2.03μs -> 2.19μs (7.31% slower)
    SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_edge_scope_is_not_set():
    # Edge: scope is not set, should ignore obj and shared_object_id only matters
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=object(), shared_object_id="id1"); result = codeflash_output # 2.12μs -> 2.33μs (9.27% slower)

def test_edge_scope_is_set_but_obj_is_missing():
    # Edge: scope is set, but obj argument is missing (None)
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}); result = codeflash_output # 2.29μs -> 787ns (191% faster)
    SHARED_OBJECT_SAVING.scope = None  # Clean up

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

def test_large_scale_many_configs():
    # Large scale: many configs in a loop
    configs = []
    for i in range(500):
        config = {"units": i, "activation": "relu"}
        codeflash_output = serialize_keras_class_and_config("Dense", config, shared_object_id=str(i)); result = codeflash_output # 326μs -> 140μs (132% faster)
        configs.append(result)
    # Ensure all configs are unique by shared_object_id
    ids = set(r[SHARED_OBJECT_KEY] for r in configs)

def test_large_scale_scope_many_objects():
    # Large scale: many objects with shared scope
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    objs = [object() for _ in range(300)]
    results = []
    for i, obj in enumerate(objs):
        config = {"units": i}
        codeflash_output = serialize_keras_class_and_config("Dense", config, obj=obj); result = codeflash_output # 191μs -> 190μs (0.918% faster)
        results.append(result)
    # Now, calling again with same obj should return same config (not create new)
    for i, obj in enumerate(objs):
        codeflash_output = serialize_keras_class_and_config("Dense", {"units": i}, obj=obj); result2 = codeflash_output # 186μs -> 183μs (1.76% faster)
    SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_large_scale_nested_config():
    # Large scale: deeply nested config
    nested_config = {"layers": [{"type": "Dense", "config": {"units": i, "activation": "relu"}} for i in range(100)]}
    codeflash_output = serialize_keras_class_and_config("Model", nested_config, shared_object_id="model_big"); result = codeflash_output # 2.25μs -> 953ns (136% faster)

def test_large_scale_scope_with_shared_object_id():
    # Large scale: scope + shared_object_id for many objects
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    objs = [object() for _ in range(100)]
    for i, obj in enumerate(objs):
        config = {"units": i}
        codeflash_output = serialize_keras_class_and_config("Dense", config, obj=obj, shared_object_id=f"id_{i}"); result = codeflash_output # 67.1μs -> 65.7μs (2.19% faster)
    SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_large_scale_config_is_large_dict():
    # Large scale: config is a large dict
    config = {f"key_{i}": i for i in range(800)}
    codeflash_output = serialize_keras_class_and_config("Dense", config); result = codeflash_output # 1.86μs -> 652ns (186% faster)
    # Ensure all keys are present
    for i in range(800):
        pass
# 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.legacy.saving.serialization import \
    serialize_keras_class_and_config

SHARED_OBJECT_KEY = "shared_object_id"
SHARED_OBJECT_SAVING = threading.local()
from keras.src.legacy.saving.serialization import \
    serialize_keras_class_and_config

# Helper class to simulate a SharedObjectSavingScope for edge case tests
class DummyScope:
    def __init__(self):
        self.configs = {}
        self.created = []

    def get_config(self, obj):
        # Return config if obj is found, else None
        return self.configs.get(obj, None)

    def create_config(self, base_config, obj):
        # Save and return a config for obj
        config = dict(base_config)
        config["dummy_scope"] = True
        self.created.append((obj, config))
        self.configs[obj] = config
        return config

# ------------------- UNIT TESTS -------------------

# BASIC TEST CASES

def test_basic_serialization_minimal():
    # Minimal config, no shared_object_id or obj
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}); result = codeflash_output # 1.84μs -> 581ns (216% faster)

def test_basic_serialization_with_shared_object_id():
    # Should include shared_object_id in output
    codeflash_output = serialize_keras_class_and_config(
        "Dense", {"units": 32}, shared_object_id="abc123"
    ); result = codeflash_output # 2.06μs -> 934ns (121% faster)

def test_basic_serialization_with_obj_none_scope_none():
    # obj is not None, but scope is None, should ignore obj
    obj = object()
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj); result = codeflash_output # 2.00μs -> 2.55μs (21.6% slower)

def test_basic_serialization_with_empty_config():
    # Empty config dict
    codeflash_output = serialize_keras_class_and_config("InputLayer", {}); result = codeflash_output # 1.67μs -> 551ns (202% faster)

def test_basic_serialization_with_various_types_in_config():
    # Config contains various types
    config = {"units": 10, "activation": "relu", "trainable": True, "rate": 0.5}
    codeflash_output = serialize_keras_class_and_config("Dropout", config); result = codeflash_output # 1.76μs -> 514ns (243% faster)

# EDGE TEST CASES

def test_edge_serialization_with_shared_object_id_none():
    # shared_object_id explicitly None, should not appear in output
    codeflash_output = serialize_keras_class_and_config(
        "Dense", {"units": 32}, shared_object_id=None
    ); result = codeflash_output # 1.96μs -> 816ns (141% faster)

def test_edge_serialization_with_empty_class_name():
    # Empty string for class name
    codeflash_output = serialize_keras_class_and_config("", {"foo": "bar"}); result = codeflash_output # 1.67μs -> 539ns (210% faster)

def test_edge_serialization_with_non_string_class_name():
    # Non-string class name (int)
    codeflash_output = serialize_keras_class_and_config(123, {"foo": "bar"}); result = codeflash_output # 1.74μs -> 534ns (226% faster)

def test_edge_serialization_with_non_dict_config():
    # Non-dict config (list)
    codeflash_output = serialize_keras_class_and_config("Dense", [1, 2, 3]); result = codeflash_output # 1.68μs -> 526ns (220% faster)

def test_edge_serialization_with_none_config():
    # None config
    codeflash_output = serialize_keras_class_and_config("Dense", None); result = codeflash_output # 1.69μs -> 529ns (220% faster)

def test_edge_serialization_with_scope_and_obj_existing_config():
    # Scope returns an existing config
    scope = DummyScope()
    obj = object()
    existing_config = {"class_name": "Dense", "config": {"units": 32}, "existing": True}
    scope.configs[obj] = existing_config
    SHARED_OBJECT_SAVING.scope = scope
    try:
        codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj); result = codeflash_output
    finally:
        SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_edge_serialization_with_scope_and_obj_create_config():
    # Scope returns None, so create_config is called
    scope = DummyScope()
    obj = object()
    SHARED_OBJECT_SAVING.scope = scope
    try:
        codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj); result = codeflash_output
    finally:
        SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_edge_serialization_with_scope_obj_and_shared_object_id():
    # Scope is active, obj is not None, shared_object_id is set
    scope = DummyScope()
    obj = object()
    SHARED_OBJECT_SAVING.scope = scope
    try:
        codeflash_output = serialize_keras_class_and_config(
            "Dense", {"units": 32}, obj=obj, shared_object_id="id42"
        ); result = codeflash_output
    finally:
        SHARED_OBJECT_SAVING.scope = None  # Clean up

def test_edge_serialization_with_scope_none_obj_not_none():
    # Scope is None, obj is not None: should ignore obj
    obj = object()
    SHARED_OBJECT_SAVING.scope = None
    codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=obj); result = codeflash_output # 1.78μs -> 1.86μs (4.08% slower)

def test_edge_serialization_with_scope_not_none_obj_none():
    # Scope is not None, obj is None: should ignore scope
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    try:
        codeflash_output = serialize_keras_class_and_config("Dense", {"units": 32}, obj=None); result = codeflash_output
    finally:
        SHARED_OBJECT_SAVING.scope = None

# LARGE SCALE TEST CASES

def test_large_scale_many_configs():
    # Test with a large number of configs
    configs = [{"units": i, "activation": "relu"} for i in range(1000)]
    for i, cfg in enumerate(configs):
        codeflash_output = serialize_keras_class_and_config("Dense", cfg); result = codeflash_output # 547μs -> 200μs (173% faster)

def test_large_scale_many_shared_object_ids():
    # Test with many unique shared_object_ids
    for i in range(1000):
        codeflash_output = serialize_keras_class_and_config(
            "Dense", {"units": i}, shared_object_id=f"id_{i}"
        ); result = codeflash_output # 634μs -> 265μs (139% faster)

def test_large_scale_many_objs_with_scope():
    # Test with many objects and a shared scope
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    try:
        objs = [object() for _ in range(1000)]
        for i, obj in enumerate(objs):
            codeflash_output = serialize_keras_class_and_config(
                "Dense", {"units": i}, obj=obj, shared_object_id=f"id_{i}"
            ); result = codeflash_output
    finally:
        SHARED_OBJECT_SAVING.scope = None

def test_large_scale_repeated_obj_with_scope():
    # Test that repeated calls with the same obj return the same config from scope
    scope = DummyScope()
    SHARED_OBJECT_SAVING.scope = scope
    try:
        obj = object()
        codeflash_output = serialize_keras_class_and_config("Dense", {"units": 42}, obj=obj); result1 = codeflash_output
        codeflash_output = serialize_keras_class_and_config("Dense", {"units": 42}, obj=obj); result2 = codeflash_output
    finally:
        SHARED_OBJECT_SAVING.scope = None

def test_large_scale_with_long_class_name_and_config():
    # Very long class name and large config dict
    class_name = "Dense" * 100
    config = {f"param_{i}": i for i in range(1000)}
    codeflash_output = serialize_keras_class_and_config(class_name, config); result = codeflash_output # 1.79μs -> 623ns (187% 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-serialize_keras_class_and_config-mjaew037 and push.