⚡️ Speed up function keras_model_summary by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for keras_model_summary in keras/src/callbacks/tensorboard.py

⏱️ Runtime : 2.84 milliseconds → 2.67 milliseconds (best of 118 runs)

📝 Explanation and details

The optimization achieves a 6% speedup by eliminating redundant imports on every function call. The key change is moving expensive TensorFlow imports to module-level - tensorflow.summary and SummaryMetadata are now imported once at startup rather than repeatedly inside the function.

What changed:

• Moved import tensorflow.summary and from tensorflow.compat.v1 import SummaryMetadata to module-level with underscore prefixes
• Added minor attribute lookup optimization by caching _summary.experimental.summary_scope and _summary.write before the context manager

Why it's faster:
The line profiler shows the original imports consumed 1.9% of total runtime (325,013 + 206,389 ns out of 27.69ms total). TensorFlow modules are notoriously expensive to import due to their complex initialization. By importing once at module load instead of on every call, we eliminate this repeated overhead entirely.

Impact on workloads:
Looking at the function reference, keras_model_summary is called from _write_keras_model_summary() in TensorBoard callbacks. While this appears to be called once per model (step=0), the optimization is particularly valuable because:

• TensorBoard callbacks are commonly used in training pipelines where even small overhead matters
• The test results show 7% speedup in the "many calls" scenario, indicating the optimization scales well when called repeatedly
• Any reduction in callback overhead directly improves training iteration time

The optimization is most effective for workloads that call this function multiple times or where import overhead is a concern, with minimal impact on single-use scenarios.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import sys
import types
# function to test
import warnings

# imports
import pytest  # used for our unit tests
from keras.src.callbacks.tensorboard import keras_model_summary

# unit tests

class DummySummaryWriter:
    """Dummy summary writer for testing purposes."""
    def __init__(self):
        self.written = []

    def write(self, tag, tensor, step, metadata):
        # Simulate writing summary, store call arguments
        self.written.append((tag, tensor, step, metadata))
        return True

class DummySummaryScope:
    """Dummy summary scope context manager."""
    def __init__(self, name, plugin_name, args):
        self.name = name
        self.plugin_name = plugin_name
        self.args = args

    def __enter__(self):
        # Return tag and dummy value
        return (self.name, None)

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

class DummySummaryExperimental:
    """Dummy experimental module for summary."""
    def __init__(self, writer):
        self.writer = writer
        self._step = None

    def summary_scope(self, name, plugin_name, args):
        return DummySummaryScope(name, plugin_name, args)

    def get_step(self):
        return self._step

class DummySummary:
    """Dummy summary module."""
    def __init__(self, writer):
        self.experimental = DummySummaryExperimental(writer)
        self.write = writer.write

class DummySummaryMetadataPluginData:
    def __init__(self):
        self.plugin_name = None
        self.content = None

class DummySummaryMetadata:
    def __init__(self):
        self.plugin_data = DummySummaryMetadataPluginData()

@pytest.fixture(autouse=True)
def patch_tf_summary(monkeypatch):
    # Patch tensorflow.summary and tensorflow.compat.v1.SummaryMetadata
    dummy_writer = DummySummaryWriter()
    dummy_summary = DummySummary(dummy_writer)
    monkeypatch.setitem(sys.modules, "tensorflow.summary", dummy_summary)
    monkeypatch.setitem(sys.modules, "tensorflow.compat.v1", types.SimpleNamespace(SummaryMetadata=DummySummaryMetadata))
    yield dummy_writer  # yield for inspection if needed

# Helper dummy keras-like model
class DummyKerasModel:
    def to_json(self):
        return '{"class_name": "DummyKerasModel", "config": {}}'

class DummyKerasModelFail:
    def to_json(self):
        raise Exception("Serialization failed")

# Basic Test Cases

def test_edge_serialization_failure(patch_tf_summary):
    """Test when model.to_json() raises an exception."""
    model = DummyKerasModelFail()
    with pytest.warns(UserWarning) as record:
        codeflash_output = keras_model_summary("fail_model", model, step=1); result = codeflash_output # 13.7μs -> 15.0μs (8.62% slower)

def test_large_scale_many_calls(patch_tf_summary):
    """Test calling keras_model_summary many times in a loop."""
    model = DummyKerasModel()
    for i in range(100):  # 100 calls, well under 1000
        codeflash_output = keras_model_summary(f"model_{i}", model, step=i); result = codeflash_output # 1.11ms -> 1.04ms (6.91% faster)

def test_large_scale_many_serialization_failures(patch_tf_summary):
    """Test many serialization failures."""
    for i in range(20):
        class BadModel:
            def to_json(self):
                raise Exception(f"Fail {i}")
        model = BadModel()
        with pytest.warns(UserWarning):
            codeflash_output = keras_model_summary(f"bad_model_{i}", model, step=i); result = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import sys
import types
# function to test
import warnings

# imports
import pytest
from keras.src.callbacks.tensorboard import keras_model_summary

# unit tests

# To avoid requiring a real TensorFlow/Keras install, we will monkeypatch
# the relevant parts of tensorflow.summary and keras.Model for testing.

class DummySummaryWriter:
    def __init__(self):
        self.written = []

    def write(self, tag, tensor, step, metadata):
        # Simulate writing and always return True
        self.written.append((tag, tensor, step, metadata))
        return True

class DummySummaryScope:
    def __init__(self, tag):
        self.tag = tag

    def __enter__(self):
        # Return (tag, None)
        return (self.tag, None)

    def __exit__(self, exc_type, exc_val, exc_tb):
        return False

class DummySummaryExperimental:
    def __init__(self, writer):
        self.writer = writer

    def summary_scope(self, name, default_name, values):
        # Always return a dummy scope with a tag
        return DummySummaryScope(f"scope_{name}")

class DummySummary:
    def __init__(self, writer):
        self.experimental = DummySummaryExperimental(writer)
        self.write = writer.write

class DummySummaryMetadataPluginData:
    def __init__(self):
        self.plugin_name = None
        self.content = None

class DummySummaryMetadata:
    def __init__(self):
        self.plugin_data = DummySummaryMetadataPluginData()

class DummyModel:
    def __init__(self, json_return="{}"):
        self._json_return = json_return
        self.to_json_called = False

    def to_json(self):
        self.to_json_called = True
        return self._json_return

class DummyModelRaises:
    def to_json(self):
        raise RuntimeError("Serialization failed.")

# Patch tensorflow.summary and tensorflow.compat.v1.SummaryMetadata for tests
@pytest.fixture(autouse=True)
def patch_tf_summary(monkeypatch):
    # Patch summary
    writer = DummySummaryWriter()
    dummy_summary = DummySummary(writer)
    monkeypatch.setitem(sys.modules, "tensorflow.summary", dummy_summary)
    # Patch SummaryMetadata
    monkeypatch.setitem(sys.modules, "tensorflow.compat.v1", types.SimpleNamespace(
        SummaryMetadata=DummySummaryMetadata
    ))
    yield writer

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

def test_basic_with_different_names(patch_tf_summary):
    # Test with different summary names
    model = DummyModel(json_return='{"foo": 123}')
    codeflash_output = keras_model_summary("foo_bar", model, step=42); result = codeflash_output # 46.3μs -> 46.9μs (1.22% slower)

def test_basic_step_none_raises(monkeypatch):
    # Test when step is None and tf.summary.experimental.get_step() is None
    # Simulate get_step returning None and default writer exists
    class DummySummaryExperimentalWithStep(DummySummaryExperimental):
        def get_step(self):
            return None
    writer = DummySummaryWriter()
    dummy_summary = DummySummary(writer)
    dummy_summary.experimental = DummySummaryExperimentalWithStep(writer)
    monkeypatch.setitem(sys.modules, "tensorflow.summary", dummy_summary)
    monkeypatch.setitem(sys.modules, "tensorflow.compat.v1", types.SimpleNamespace(
        SummaryMetadata=DummySummaryMetadata
    ))
    model = DummyModel(json_return='{"bar": 1}')
    # Should raise ValueError if step is None and get_step is None
    # But our dummy implementation does not call get_step, so we skip this test.
    # (The actual code only raises if there's a default writer and no step.)
    # Here, we just check that step=None works if not raising.
    codeflash_output = keras_model_summary("bar", model, step=None); result = codeflash_output # 35.8μs -> 35.5μs (0.902% faster)

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

def test_model_to_json_raises_returns_false(monkeypatch):
    # Test when model.to_json raises: should return False and warn
    model = DummyModelRaises()
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = keras_model_summary("broken", model, step=5); result = codeflash_output # 13.1μs -> 13.0μs (0.679% faster)

def test_non_string_json(monkeypatch):
    # Test if model.to_json returns a non-string (should still be passed to write)
    model = DummyModel(json_return=12345)
    codeflash_output = keras_model_summary("non_string", model, step=2); result = codeflash_output # 37.5μs -> 36.6μs (2.22% faster)

def test_empty_name(monkeypatch):
    # Test with empty string for name
    model = DummyModel(json_return='{"empty": true}')
    codeflash_output = keras_model_summary("", model, step=10); result = codeflash_output # 33.4μs -> 32.2μs (3.88% faster)

def test_empty_json(monkeypatch):
    # Test with model.to_json returning empty string
    model = DummyModel(json_return="")
    codeflash_output = keras_model_summary("empty_json", model, step=3); result = codeflash_output # 32.7μs -> 33.2μs (1.71% slower)

def test_step_zero(monkeypatch):
    # Test with step=0
    model = DummyModel(json_return='{"zero": 0}')
    codeflash_output = keras_model_summary("zero_step", model, step=0); result = codeflash_output # 32.7μs -> 32.6μs (0.556% faster)

def test_name_with_special_chars(monkeypatch):
    # Test with name containing special characters
    model = DummyModel(json_return='{"special": "yes"}')
    codeflash_output = keras_model_summary("special!@# $%^&*()", model, step=7); result = codeflash_output # 34.3μs -> 33.3μs (3.23% faster)

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

def test_large_json(monkeypatch):
    # Test with a large JSON string (but <100MB)
    large_json = '{"data": "' + ("x" * 10**6) + '"}'  # ~1MB string
    model = DummyModel(json_return=large_json)
    codeflash_output = keras_model_summary("large", model, step=100); result = codeflash_output # 48.7μs -> 47.4μs (2.88% faster)

def test_many_calls(monkeypatch):
    # Call keras_model_summary many times in a loop (scalability)
    for i in range(100):
        model = DummyModel(json_return=f'{{"idx": {i}}}')
        codeflash_output = keras_model_summary(f"loop_{i}", model, step=i); result = codeflash_output # 1.09ms -> 1.02ms (7.18% faster)

def test_large_number_of_layers(monkeypatch):
    # Simulate a model with a very large JSON (representing many layers)
    layer_json = ",".join([f'"layer{i}":{{"type":"Dense"}}' for i in range(300)])
    model_json = "{" + layer_json + "}"
    model = DummyModel(json_return=model_json)
    codeflash_output = keras_model_summary("many_layers", model, step=123); result = codeflash_output # 32.4μs -> 31.7μs (2.16% faster)

def test_different_types_for_step(monkeypatch):
    # Test with step as different int types (simulate int64-castable)
    model = DummyModel(json_return='{"foo": "bar"}')
    for s in [1, 1_000_000_000_000, 0]:
        codeflash_output = keras_model_summary("step_type", model, step=s); result = codeflash_output # 59.3μs -> 58.8μs (0.852% faster)

def test_multiple_models(monkeypatch):
    # Test calling with different model instances
    models = [DummyModel(json_return=f'{{"id": {i}}}') for i in range(10)]
    for i, model in enumerate(models):
        codeflash_output = keras_model_summary(f"multi_{i}", model, step=i); result = codeflash_output # 135μs -> 132μs (2.35% 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-keras_model_summary-mjaatxy7 and push.