Add none to aggregation

keras/src/backend/common/variables.py

@@ -95,10 +95,10 @@ def __init__(
                 "cannot contain character `/`. "
                 f"Received: name={name}"
             )
-        if aggregation not in ("mean", "sum", "only_first_replica"):
+        if aggregation not in ("none", "mean", "sum", "only_first_replica"):
             raise ValueError(
                 "Invalid valid for argument `aggregation`. Expected "
-                "one of {'mean', 'sum', 'only_first_replica'}. "
+                "one of {'none', 'mean', 'sum', 'only_first_replica'}. "
                 f"Received: aggregation={aggregation}"
             )
         self.name = name

keras/src/backend/tensorflow/core.py

@@ -120,11 +120,12 @@ def _write_object_proto(self, proto, options):
 
     def _map_aggregation(self, aggregation):
         mapping = {
+            "none": tf.VariableAggregation.NONE,
             "sum": tf.VariableAggregation.SUM,
             "mean": tf.VariableAggregation.MEAN,
             "only_first_replica": tf.VariableAggregation.ONLY_FIRST_REPLICA,
         }
-        return mapping.get(aggregation, tf.VariableAggregation.NONE)
+        return mapping[aggregation]
 
 
 def convert_to_tensor(x, dtype=None, sparse=None):

keras/src/optimizers/loss_scale_optimizer.py

@@ -67,12 +67,14 @@ def build(self, var_list):
             shape=(),
             dtype="int",
             initializer=initializers.Zeros(),
+            aggregation="none",
             name="step_counter",
         )
         self.dynamic_scale = self.add_variable(
             shape=(),
             dtype="float32",
             initializer=initializers.Constant(self.initial_scale),
+            aggregation="none",
             name="dynamic_scale",
         )
         self.inner_optimizer.build(var_list)