Add rng_type to BaseRandomLayer.get_config().

The stateful mode can be enabled either by passing it explicitly to the layer's `__init__` or by using `enable_tf_random_generator`. However, it was not saved in either cases.

PiperOrigin-RevId: 775020954

Author: hertschuh

tf_keras/engine/base_layer.py

@@ -3932,6 +3932,17 @@ def build(self, input_shape):
         super().build(input_shape)
         self._random_generator._maybe_init()
 
+    def get_config(self):
+        base_config = super().get_config()
+        if (
+            self._random_generator._rng_type
+            == backend.RandomGenerator.RNG_LEGACY_STATEFUL
+        ):
+            return base_config
+
+        config = {"rng_type": self._random_generator._rng_type}
+        return dict(list(base_config.items()) + list(config.items()))
+
     def _trackable_children(self, save_type="checkpoint", **kwargs):
         if save_type == "savedmodel":
             cache = kwargs["cache"]

tf_keras/engine/base_layer_test.py

@@ -1272,6 +1272,57 @@ def call(self, inputs):
         self.assertAllEqual(out_false, sample_input)
 
 
+@test_utils.run_v2_only
+class BaseRandomLayerTest(test_combinations.TestCase):
+    def teardown(self):
+        backend.disable_tf_random_generator()
+
+    def test_rng_type_is_saved_in_config(self):
+        backend.disable_tf_random_generator()
+
+        layer = base_layer.BaseRandomLayer(rng_type="stateful")
+        config = layer.get_config()
+        self.assertEqual(config["rng_type"], "stateful")
+        reloaded_layer = base_layer.BaseRandomLayer.from_config(config)
+        self.assertEqual(reloaded_layer._random_generator._rng_type, "stateful")
+
+        layer = base_layer.BaseRandomLayer(rng_type="stateless")
+        config = layer.get_config()
+        self.assertEqual(config["rng_type"], "stateless")
+        reloaded_layer = base_layer.BaseRandomLayer.from_config(config)
+        self.assertEqual(
+            reloaded_layer._random_generator._rng_type, "stateless"
+        )
+
+        layer = base_layer.BaseRandomLayer()
+        config = layer.get_config()
+        self.assertNotIn("rng_type", config)
+        reloaded_layer = base_layer.BaseRandomLayer.from_config(config)
+        self.assertEqual(
+            reloaded_layer._random_generator._rng_type, "legacy_stateful"
+        )
+
+        layer = base_layer.BaseRandomLayer(rng_type="legacy_stateful")
+        config = layer.get_config()
+        self.assertNotIn("rng_type", config)
+        reloaded_layer = base_layer.BaseRandomLayer.from_config(config)
+        self.assertEqual(
+            reloaded_layer._random_generator._rng_type, "legacy_stateful"
+        )
+
+    def test_rng_type_with_tf_random_generator(self):
+        # Test `rng_type` is still serialized when global stateful mode is on.
+        backend.enable_tf_random_generator()
+
+        layer = base_layer.BaseRandomLayer()
+        config = layer.get_config()
+        self.assertEqual(config["rng_type"], "stateful")
+
+        backend.disable_tf_random_generator()
+        reloaded_layer = base_layer.BaseRandomLayer.from_config(config)
+        self.assertEqual(reloaded_layer._random_generator._rng_type, "stateful")
+
+
 @test_utils.run_v2_only
 class SymbolicSupportTest(test_combinations.TestCase):
     def test_using_symbolic_tensors_with_tf_ops(self):