Fix SeedGenerator in tf.distribute.MirroredStrategy (#20612)

Author: james77777778

keras/src/backend/tensorflow/core.py

@@ -44,13 +44,7 @@ def _initialize(self, value):
         )
 
     def _initialize_with_initializer(self, initializer):
-        self._value = tf.Variable(
-            lambda: initializer(self._shape, dtype=self._dtype),
-            dtype=self._dtype,
-            trainable=self.trainable,
-            name=self.name,
-            aggregation=self._map_aggregation(self.aggregation),
-        )
+        self._initialize(lambda: initializer(self._shape, dtype=self._dtype))
 
     def _deferred_initialize(self):
         if self._value is not None:

keras/src/backend/tensorflow/distribute_test.py

@@ -137,6 +137,14 @@ def test_variable_aggregation(self):
             self.assertEqual(v2.aggregation, "sum")
             self.assertEqual(v2.value.aggregation, tf.VariableAggregation.SUM)
 
+    def test_seed_generator(self):
+        strategy = tf.distribute.MirroredStrategy(["CPU:0", "CPU:1"])
+        with strategy.scope():
+            seed_generator = keras.random.SeedGenerator(42)
+            states = strategy.run(lambda: seed_generator.state.value).values
+            for s in states:
+                self.assertAllClose(keras.ops.convert_to_numpy(s), (42, 0))
+
     def test_correctness_with_fit_and_regularizer(self):
         strategy = tf.distribute.MirroredStrategy(["CPU:0", "CPU:1"])
 

keras/src/random/seed_generator.py

@@ -89,6 +89,7 @@ def seed_initializer(*args, **kwargs):
                 shape=(2,),
                 dtype=self.backend.random_seed_dtype(),
                 trainable=False,
+                aggregation="none",
                 name="seed_generator_state",
             )