def operation_fn(operation, **call_context_args):
    """Wraps each op to inject the call-context args."""

    def call(*args, **kwargs):
        # Propagate all registered call-context args
        for name, value in call_context_args.items():
            if (
                name in getattr(operation, "_call_context_args", {})
                and value is not None
            ):
                kwargs[name] = value

        return operation(*args, **kwargs)

    return call
Wraps each op to inject the call-context args.
keras/src/models/functional.py