def test_register_call_context_arguments(self):
        """Validate that registering call-context args works as expected."""

        class MyLayer(layers.Layer):
            def call(self, x):
                return x

        layer = MyLayer()

        layer._register_call_context_args("foo_mode")

        self.assertCountEqual(
            layer._call_context_args, ("foo_mode", "training")
        )
Validate that registering call-context args works as expected.
keras/src/layers/layer_test.py