Add custom variable updater.

Allows customization for how variables are updated by the optimizer.
The base optimizer simply defers to the update handler to do the
update, allowing full customization.

Can replace the existing `overwrite_with_gradient` attribute on
variables, which currently is very application-specific.

Eliminates creation of optimizer variables that have custom updaters
(including `overwrite_with_gradient`), since those variables are
never used and may be wasteful.

Author: cantonios

keras/src/backend/common/variables.py

@@ -150,6 +150,8 @@ def __init__(
         self._autocast = bool(autocast)
         self._aggregation = aggregation
         self._synchronization = synchronization
+        # Custom variable updater.
+        self._updater = None
         # `self._overwrite_with_gradient` is an internal property to determine
         # whether this variable should be overwritten by the computed gradient.
         # Ref: https://github.com/google/flax/blob/main/flax/linen/fp8_ops.py
@@ -334,6 +336,29 @@ def path(self):
         """The path of the variable within the Keras model or layer."""
         return self._path
 
+    @property
+    def updater(self):
+        """Custom variable updater.
+
+        This property is designed for special-casing variable updates during
+        training, such as quantized float8 `scale` and `amax_history`, where
+        the gradients represent updated scale factors, or for updating large
+        embedding tables, where we need to handle sparse updates to a dense
+        table.
+        """
+        return self._updater
+
+    @updater.setter
+    def updater(self, updater):
+        from keras.src import optimizers
+
+        if not isinstance(updater, optimizers.VariableUpdater):
+            raise TypeError(
+                "`updater` must be a `keras.optimizers.VariableUpdater`. "
+                f"Received: {updater.__class__.__name__}."
+            )
+        self._updater = updater
+
     @property
     def overwrite_with_gradient(self):
         """Whether this variable should be overwritten by the gradient.
@@ -355,6 +380,9 @@ def overwrite_with_gradient(self, value):
                 f"Received: {value}"
             )
         self._overwrite_with_gradient = value
+        from keras.src import optimizers
+
+        self._updater = optimizers.OverwriteScaleWithGradientUpdater()
 
     @property
     def regularizer(self):

keras/src/optimizers/__init__.py

@@ -11,6 +11,8 @@
 from keras.src.optimizers.muon import Muon
 from keras.src.optimizers.nadam import Nadam
 from keras.src.optimizers.optimizer import Optimizer
+from keras.src.optimizers.optimizer import OverwriteScaleWithGradientUpdater
+from keras.src.optimizers.optimizer import VariableUpdater
 from keras.src.optimizers.rmsprop import RMSprop
 from keras.src.optimizers.sgd import SGD
 from keras.src.saving import serialization_lib

keras/src/optimizers/adadelta.py

@@ -75,15 +75,12 @@ def build(self, var_list):
         if self.built:
             return
         super().build(var_list)
-        self._accumulated_grads = []
-        self._accumulated_delta_vars = []
-        for var in var_list:
-            self._accumulated_grads.append(
-                self.add_variable_from_reference(var, "accumulated_grad")
-            )
-            self._accumulated_delta_vars.append(
-                self.add_variable_from_reference(var, "accumulated_delta_var")
-            )
+        self._accumulated_grads = self.add_optimizer_variables(
+            var_list, "accumulated_grad"
+        )
+        self._accumulated_delta_vars = self.add_optimizer_variables(
+            var_list, "accumulated_delta_var"
+        )
 
     def update_step(self, grad, variable, learning_rate):
         """Update step given gradient and the associated model variable."""

keras/src/optimizers/adafactor.py

@@ -1,4 +1,3 @@
-from keras.src import backend
 from keras.src import ops
 from keras.src.api_export import keras_export
 from keras.src.optimizers import optimizer
@@ -97,16 +96,11 @@ def build(self, var_list):
         self._c = []
         self._v = []
         for var in var_list:
-            if len(var.shape) < 2:
-                # Don't factor if variable is of dimension < 2, but we still
-                # need to create dummy variables as placeholder.
-                with backend.name_scope(self.name, caller=self):
-                    self._r.append(
-                        backend.Variable(0, name=var.name, trainable=False)
-                    )
-                    self._c.append(
-                        backend.Variable(0, name=var.name, trainable=False)
-                    )
+            variable_updater = self._get_variable_updater(var)
+            if len(var.shape) < 2 or variable_updater is not None:
+                # Don't factor if variable is of dimension < 2.
+                self._r.append(None)
+                self._c.append(None)
             else:
                 # Always factor the last 2 dimensions.
                 r_shape = var.shape[:-1]
@@ -125,11 +119,15 @@ def build(self, var_list):
                         name=var.name,
                     )
                 )
-            self._v.append(
-                self.add_variable_from_reference(
-                    reference_variable=var, name="velocity"
+
+            if variable_updater is not None:
+                self._v.append(None)
+            else:
+                self._v.append(
+                    self.add_variable_from_reference(
+                        reference_variable=var, name="velocity"
+                    )
                 )
-            )
 
     def _rms(self, x):
         return ops.sqrt(ops.mean(ops.square(x)))

keras/src/optimizers/adagrad.py

@@ -70,17 +70,10 @@ def build(self, var_list):
         if self.built:
             return
         super().build(var_list)
-        self._accumulators = []
         initializer = initializers.Constant(self.initial_accumulator_value)
-        for var in var_list:
-            self._accumulators.append(
-                self.add_variable(
-                    shape=var.shape,
-                    initializer=initializer,
-                    dtype=var.dtype,
-                    name="accumulator",
-                )
-            )
+        self._accumulators = self.add_optimizer_variables(
+            var_list, "accumulator", initializer=initializer
+        )
 
     def update_step(self, gradient, variable, learning_rate):
         """Update step given gradient and the associated model variable."""

keras/src/optimizers/adam.py

@@ -90,27 +90,13 @@ def build(self, var_list):
         if self.built:
             return
         super().build(var_list)
-        self._momentums = []
-        self._velocities = []
-        for var in var_list:
-            self._momentums.append(
-                self.add_variable_from_reference(
-                    reference_variable=var, name="momentum"
-                )
-            )
-            self._velocities.append(
-                self.add_variable_from_reference(
-                    reference_variable=var, name="velocity"
-                )
-            )
+        self._momentums = self.add_optimizer_variables(var_list, "momentum")
+        self._velocities = self.add_optimizer_variables(var_list, "velocity")
+
         if self.amsgrad:
-            self._velocity_hats = []
-            for var in var_list:
-                self._velocity_hats.append(
-                    self.add_variable_from_reference(
-                        reference_variable=var, name="velocity_hat"
-                    )
-                )
+            self._velocity_hats = self.add_optimizer_variables(
+                var_list, "velocity_hat"
+            )
 
     def update_step(self, gradient, variable, learning_rate):
         """Update step given gradient and the associated model variable."""

keras/src/optimizers/adamax.py

@@ -98,19 +98,8 @@ def build(self, var_list):
         if self.built:
             return
         super().build(var_list)
-        self._m = []
-        self._u = []
-        for var in var_list:
-            self._m.append(
-                self.add_variable_from_reference(
-                    reference_variable=var, name="momentum"
-                )
-            )
-            self._u.append(
-                self.add_variable_from_reference(
-                    reference_variable=var, name="norm"
-                )
-            )
+        self._m = self.add_optimizer_variables(var_list, "momentum")
+        self._u = self.add_optimizer_variables(var_list, "norm")
 
     def update_step(self, gradient, variable, learning_rate):
         """Update step given gradient and the associated model variable."""