wraps arbitrary expressions as a layer object the lambda layer exists so that arbitrary expressions can be used as a layer when constructing sequential and functional api models lambda layers are best suited for simple operations or quick experimentation for more advanced use cases prefer writing new subclasses of layer warning lambda layers have deserialization limitations the main reason to subclass layer instead of using a lambda layer is saving and inspecting a model lambda layers are saved by serializing the python bytecode which is fundamentally nonportable and potentially unsafe they should only be loaded in the same environment where they were saved subclassed layers can be saved in a more portable way by overriding their getconfig method models that rely on subclassed layers are also often easier to visualize and reason about example python add a x x2 layer model addlambdalambda x x 2 args function the function to be evaluated takes input tensor as first argument outputshape expected output shape from function this argument can usually be inferred if not explicitly provided can be a tuple or function if a tuple it only specifies the first dimension onward sample dimension is assumed either the same as the input outputshape inputshape0 outputshape or the input is none and the sample dimension is also none outputshape none outputshape if a function it specifies the entire shape as a function of the input shape outputshape finputshape mask either none indicating no masking or a callable with the same signature as the computemask layer method or a tensor that will be returned as output mask regardless of what the input is arguments optional dictionary of keyword arguments to be passed to the function warning on every invocation will be quite irksome in eager mode leverage backend shape inference output shapes are passed directly and don t include batch dimension we must copy for thread safety but it only needs to be a shallow copy wraps arbitrary expressions as a layer object the lambda layer exists so that arbitrary expressions can be used as a layer when constructing sequential and functional api models lambda layers are best suited for simple operations or quick experimentation for more advanced use cases prefer writing new subclasses of layer warning lambda layers have de serialization limitations the main reason to subclass layer instead of using a lambda layer is saving and inspecting a model lambda layers are saved by serializing the python bytecode which is fundamentally non portable and potentially unsafe they should only be loaded in the same environment where they were saved subclassed layers can be saved in a more portable way by overriding their get_config method models that rely on subclassed layers are also often easier to visualize and reason about example python add a x x 2 layer model add lambda lambda x x 2 args function the function to be evaluated takes input tensor as first argument output_shape expected output shape from function this argument can usually be inferred if not explicitly provided can be a tuple or function if a tuple it only specifies the first dimension onward sample dimension is assumed either the same as the input output_shape input_shape 0 output_shape or the input is none and the sample dimension is also none output_shape none output_shape if a function it specifies the entire shape as a function of the input shape output_shape f input_shape mask either none indicating no masking or a callable with the same signature as the compute_mask layer method or a tensor that will be returned as output mask regardless of what the input is arguments optional dictionary of keyword arguments to be passed to the function warning on every invocation will be quite irksome in eager mode leverage backend shape inference output shapes are passed directly and don t include batch dimension we must copy for thread safety but it only needs to be a shallow copy
import inspect
import types

import tree

from keras import backend
from keras.api_export import keras_export
from keras.layers.layer import Layer
from keras.saving import serialization_lib
from keras.utils import python_utils
from keras.utils import shape_utils


@keras_export("keras.layers.Lambda")
class Lambda(Layer):
    

    def __init__(
        self, function, output_shape=None, mask=None, arguments=None, **kwargs
    ):
        super().__init__(**kwargs)

        self.arguments = arguments or {}
        self.function = function

        if mask is not None:
            self.supports_masking = True
        else:
            self.supports_masking = False
        self.mask = mask
        self._output_shape = output_shape

        
        self._already_warned = False

        function_args = inspect.getfullargspec(function).args
        self._fn_expects_training_arg = "training" in function_args
        self._fn_expects_mask_arg = "mask" in function_args

    def compute_output_shape(self, input_shape):
        if self._output_shape is None:
            
            try:
                inputs = shape_utils.map_shape_structure(
                    lambda x: backend.KerasTensor(x, dtype=self.compute_dtype),
                    input_shape,
                )
                output_spec = backend.compute_output_spec(self.call, inputs)
                return tree.map_structure(lambda x: x.shape, output_spec)
            except:
                raise NotImplementedError(
                    "We could not automatically infer the shape of "
                    "the Lambda's output. Please specify the `output_shape` "
                    "argument for this Lambda layer."
                )

        if callable(self._output_shape):
            return self._output_shape(input_shape)

        
        batch_size = tree.flatten(input_shape)[0]

        def _add_batch(shape):
            return (batch_size,) + shape

        return shape_utils.map_shape_structure(_add_batch, self._output_shape)

    def call(self, inputs, mask=None, training=None):
        
        
        kwargs = {k: v for k, v in self.arguments.items()}
        if self._fn_expects_mask_arg:
            kwargs["mask"] = mask
        if self._fn_expects_training_arg:
            kwargs["training"] = training
        return self.function(inputs, **kwargs)

    def compute_mask(self, inputs, mask=None):
        if callable(self.mask):
            return self.mask(inputs, mask)
        return self.mask

    def get_config(self):
        config = {
            "function": self._serialize_function_to_config(self.function),
        }
        if self._output_shape is not None:
            if callable(self._output_shape):
                output_shape = self._serialize_function_to_config(
                    self._output_shape
                )
            else:
                output_shape = self._output_shape
            config["output_shape"] = output_shape
        if self.mask is not None:
            if callable(self.mask):
                mask = self._serialize_function_to_config(self.mask)
            else:
                mask = serialization_lib.serialize_keras_object(self.mask)
            config["mask"] = mask
        config["arguments"] = serialization_lib.serialize_keras_object(
            self.arguments
        )
        base_config = super().get_config()
        return {**base_config, **config}

    def _serialize_function_to_config(self, fn):
        if isinstance(fn, types.LambdaType) and fn.__name__ == "<lambda>":
            code, defaults, closure = python_utils.func_dump(fn)
            return {
                "class_name": "__lambda__",
                "config": {
                    "code": code,
                    "defaults": defaults,
                    "closure": closure,
                },
            }
        elif callable(fn):
            return serialization_lib.serialize_keras_object(fn)
        raise ValueError(
            "Invalid input type for serialization. "
            f"Received: {fn} of type {type(fn)}."
        )

    @staticmethod
    def _raise_for_lambda_deserialization(arg_name, safe_mode):
        if safe_mode:
            raise ValueError(
                "The `{arg_name}` of this `Lambda` layer is a Python lambda. "
                "Deserializing it is unsafe. If you trust the source of the "
                "config artifact, you can override this error "
                "by passing `safe_mode=False` "
                "to `from_config()`, or calling "
                "`keras.config.enable_unsafe_deserialization()."
            )

    @classmethod
    def from_config(cls, config, custom_objects=None, safe_mode=None):
        safe_mode = safe_mode or serialization_lib.in_safe_mode()
        fn_config = config["function"]
        if (
            isinstance(fn_config, dict)
            and "class_name" in fn_config
            and fn_config["class_name"] == "__lambda__"
        ):
            cls._raise_for_lambda_deserialization("function", safe_mode)
            inner_config = fn_config["config"]
            fn = python_utils.func_load(
                inner_config["code"],
                defaults=inner_config["defaults"],
                closure=inner_config["closure"],
            )
            config["function"] = fn
        else:
            config["function"] = serialization_lib.deserialize_keras_object(
                fn_config, custom_objects=custom_objects
            )
        if "output_shape" in config:
            fn_config = config["output_shape"]
            if (
                isinstance(fn_config, dict)
                and "class_name" in fn_config
                and fn_config["class_name"] == "__lambda__"
            ):
                cls._raise_for_lambda_deserialization("function", safe_mode)
                inner_config = fn_config["config"]
                fn = python_utils.func_load(
                    inner_config["code"],
                    defaults=inner_config["defaults"],
                    closure=inner_config["closure"],
                )
                config["output_shape"] = fn
            else:
                config[
                    "output_shape"
                ] = serialization_lib.deserialize_keras_object(
                    fn_config, custom_objects=custom_objects
                )
        if "arguments" in config:
            config["arguments"] = serialization_lib.deserialize_keras_object(
                config["arguments"], custom_objects=custom_objects
            )
        return cls(**config)