simple wrapper subclass def callself inputs kwargs return ops castself layerinputs kwargs self computedtype class wrappertesttesting testcase pytest mark requirestrainablebackend def testwrapperbasicsself self runlayertest examplewrapper initkwargs layer layers dense2 inputshape2 3 expectedoutputshape2 2 expectednumtrainableweights2 expectednumnontrainableweights0 expectednumseedgenerators0 expectednumlosses0 supportsmaskingfalse self runlayertest examplewrapper initkwargs layer layers dense2 activityregularizerl2 inputshape2 3 expectedoutputshape2 2 expectednumtrainableweights2 expectednumnontrainableweights0 expectednumseedgenerators0 expectednumlosses1 supportsmaskingfalse self runlayertest examplewrapper initkwargs layer layers dense2 activityregularizer l2 inputshape2 3 expectedoutputshape2 2 expectednumtrainableweights2 expectednumnontrainableweights0 expectednumseedgenerators0 expectednumlosses1 supportsmaskingfalse self runlayertest examplewrapper initkwargs layer layers batchnormalization inputshape2 3 expectedoutputshape2 3 expectednumtrainableweights2 expectednumnontrainableweights2 expectednumseedgenerators0 expectednumlosses0 supportsmaskingfalse def testwrapperinvalidlayerself invalidlayer this is not a valid keras layer with self assertraisesregex valueerror layer supplied to wrapper isn t a supported layer type please ensure wrapped layer is a valid keras layer layers wrapperinvalidlayer simple wrapper subclass
import pytest

from keras import layers
from keras import ops
from keras import testing


class ExampleWrapper(layers.Wrapper):
    

    def call(self, inputs, **kwargs):
        return ops.cast(self.layer(inputs, **kwargs), self.compute_dtype)


class WrapperTest(testing.TestCase):
    @pytest.mark.requires_trainable_backend
    def test_wrapper_basics(self):
        self.run_layer_test(
            ExampleWrapper,
            init_kwargs={
                "layer": layers.Dense(2),
            },
            input_shape=(2, 3),
            expected_output_shape=(2, 2),
            expected_num_trainable_weights=2,
            expected_num_non_trainable_weights=0,
            expected_num_seed_generators=0,
            expected_num_losses=0,
            supports_masking=False,
        )
        self.run_layer_test(
            ExampleWrapper,
            init_kwargs={
                "layer": layers.Dense(2, activity_regularizer="l2"),
            },
            input_shape=(2, 3),
            expected_output_shape=(2, 2),
            expected_num_trainable_weights=2,
            expected_num_non_trainable_weights=0,
            expected_num_seed_generators=0,
            expected_num_losses=1,
            supports_masking=False,
        )
        self.run_layer_test(
            ExampleWrapper,
            init_kwargs={
                "layer": layers.Dense(2),
                "activity_regularizer": "l2",
            },
            input_shape=(2, 3),
            expected_output_shape=(2, 2),
            expected_num_trainable_weights=2,
            expected_num_non_trainable_weights=0,
            expected_num_seed_generators=0,
            expected_num_losses=1,
            supports_masking=False,
        )
        self.run_layer_test(
            ExampleWrapper,
            init_kwargs={
                "layer": layers.BatchNormalization(),
            },
            input_shape=(2, 3),
            expected_output_shape=(2, 3),
            expected_num_trainable_weights=2,
            expected_num_non_trainable_weights=2,
            expected_num_seed_generators=0,
            expected_num_losses=0,
            supports_masking=False,
        )

    def test_wrapper_invalid_layer(self):
        invalid_layer = "This is not a valid Keras layer."

        with self.assertRaisesRegex(
            ValueError,
            "Layer .* supplied to Wrapper isn't a supported layer type. "
            "Please ensure wrapped layer is a valid Keras layer.",
        ):
            layers.Wrapper(invalid_layer)