keras-nightly 3.12.0.dev2025092403__py3-none-any.whl → 3.14.0.dev2026010104__py3-none-any.whl

keras/src/optimizers/adafactor.py

@@ -158,33 +158,52 @@ class Adafactor(optimizer.Optimizer):
         rho_t = ops.minimum(lr, 1 / ops.sqrt(local_step))
         alpha_t = ops.maximum(epsilon_2, self._rms(variable)) * rho_t
         regulated_grad_square = ops.add(ops.square(gradient), self.epsilon_1)
-        beta_2_t = 1 - ops.power(local_step, self.beta_2_decay)
+        beta_2_t = ops.subtract(1, ops.power(local_step, self.beta_2_decay))
 
         if len(variable.shape) >= 2:
             # `r` deletes the last dimension of gradient, so it is of shape
             # `gradient.shape[:-1]`.
             self.assign(
                 r,
-                beta_2_t * r
-                + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-1),
+                ops.add(
+                    ops.multiply(beta_2_t, r),
+                    ops.multiply(
+                        ops.subtract(1, beta_2_t),
+                        ops.mean(regulated_grad_square, axis=-1),
+                    ),
+                ),
             )
             # `c` deletes the second last dimension of gradient, so it is of
             # shape `gradient.shape[:-2] + gradient.shape[-1]`.
             self.assign(
                 c,
-                beta_2_t * c
-                + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-2),
+                ops.add(
+                    ops.multiply(beta_2_t, c),
+                    ops.multiply(
+                        ops.subtract(1, beta_2_t),
+                        ops.mean(regulated_grad_square, axis=-2),
+                    ),
+                ),
             )
             self.assign(
                 v,
-                ops.expand_dims(
-                    r / ops.mean(r, axis=-1, keepdims=True), axis=-1
-                )
-                * ops.expand_dims(c, -2),
+                ops.multiply(
+                    ops.expand_dims(
+                        ops.divide(r, ops.mean(r, axis=-1, keepdims=True)),
+                        axis=-1,
+                    ),
+                    ops.expand_dims(c, -2),
+                ),
             )
         else:
             self.assign(
-                v, beta_2_t * v + (1 - beta_2_t) * regulated_grad_square
+                v,
+                ops.add(
+                    ops.multiply(beta_2_t, v),
+                    ops.multiply(
+                        ops.subtract(1, beta_2_t), regulated_grad_square
+                    ),
+                ),
             )
 
         u_t = ops.divide(gradient, ops.sqrt(v))

keras/src/optimizers/base_optimizer.py

@@ -631,6 +631,20 @@ class BaseOptimizer(KerasSaveable):
             g_acc.assign(n_g_acc)
 
     def stateless_apply(self, optimizer_variables, grads, trainable_variables):
+        """Stateless version of `apply` that returns modified variables.
+
+        Args:
+            optimizer_variables: list of tensors containing the current values
+                for the optimizer variables. These are native tensors and not
+                `keras.Variable`s.
+            grads: list of gradients to apply.
+            trainable_variables: list of tensors containing the current values
+                for the model variables. These are native tensors and not
+                `keras.Variable`s.
+
+        Returns: A tuple containing two list of tensors, the updated
+            `trainable_variables` and the updated `optimizer_variables`.
+        """
         self._check_super_called()
 
         if not self.built:
@@ -969,10 +983,15 @@ class BaseOptimizer(KerasSaveable):
             ):
                 if average is not None:
                     not_first_step = ops.not_equal(self.iterations, 0)
-                    momentum = (
-                        ops.cast(not_first_step, var.dtype) * self.ema_momentum
+                    momentum = ops.multiply(
+                        ops.cast(not_first_step, var.dtype), self.ema_momentum
+                    )
+                    average.assign(
+                        ops.add(
+                            ops.multiply(momentum, average),
+                            ops.multiply(ops.subtract(1, momentum), var),
+                        )
                     )
-                    average.assign(momentum * average + (1 - momentum) * var)
 
     def _overwrite_model_variables_with_average_value(
         self, trainable_variables

keras/src/optimizers/loss_scale_optimizer.py

@@ -48,6 +48,7 @@ class LossScaleOptimizer(optimizer.Optimizer):
         inner_optimizer,
         initial_scale=2.0**15,
         dynamic_growth_steps=2000,
+        name=None,
         **kwargs,
     ):
         if not kwargs.pop("dynamic", True):
@@ -56,7 +57,42 @@ class LossScaleOptimizer(optimizer.Optimizer):
                 "Instead, simply set `loss_scale_factor` directly on the "
                 "`inner_optimizer`."
             )
-        super().__init__(learning_rate=0.0, **kwargs)
+
+        # Backwards compatibility code for deserialization.
+        # LossScaleOptimizer used to return all these parameters in `get_config`
+        # from `super.get_config` even though they are all non-functional. We
+        # no longer let user set them, but we have to allow the default values
+        # to be passed during deserialization to support older models.
+        base_optimizer_defaults = {
+            "weight_decay": None,
+            "clipnorm": None,
+            "global_clipnorm": None,
+            "clipvalue": None,
+            "use_ema": False,
+            "ema_momentum": 0.99,
+            "ema_overwrite_frequency": None,
+            "loss_scale_factor": None,
+            "gradient_accumulation_steps": None,
+        }
+        for arg_name, default_value in base_optimizer_defaults.items():
+            if arg_name not in kwargs:
+                continue
+            arg_value = kwargs.pop(arg_name)
+            if (
+                default_value is None and arg_value is not None
+            ) or arg_value != default_value:
+                raise ValueError(
+                    f"LossScaleOptimizer does not support `{arg_name}`. "
+                    f"Instead, set `{arg_name}` on the `inner_optimizer`."
+                )
+
+        if kwargs:
+            raise ValueError(
+                "LossScaleOptimizer does not support arguments: "
+                f"`{'`, `'.join(kwargs.keys())}`."
+            )
+
+        super().__init__(learning_rate=0.0, name=name)
         self.inner_optimizer = inner_optimizer
         self.initial_scale = initial_scale
         self.dynamic_growth_steps = dynamic_growth_steps
@@ -81,7 +117,7 @@ class LossScaleOptimizer(optimizer.Optimizer):
             name="dynamic_scale",
         )
         self.inner_optimizer.build(var_list)
-        self.built = True
+        super().build(var_list)
 
     @property
     def variables(self):
@@ -112,7 +148,7 @@ class LossScaleOptimizer(optimizer.Optimizer):
             mapping = list(zip(self.variables, optimizer_variables))
             with backend.StatelessScope(state_mapping=mapping) as scope:
                 self.step_counter.assign(0)
-                self.dynamic_scale.assign(self.dynamic_scale * 2.0)
+                self.dynamic_scale.assign(ops.multiply(self.dynamic_scale, 2.0))
             return [scope.get_current_value(v) for v in self._variables]
 
         def increment():
@@ -136,7 +172,7 @@ class LossScaleOptimizer(optimizer.Optimizer):
                 g
                 if g is None or self._overwrite_variable_with_gradient(v)
                 else ops.divide(g, scale)
-                for g, v in zip(grads, trainable_variables)
+                for g, v in zip(grads, self._trainable_variables)
             ]
             (
                 new_trainable_variables,
@@ -156,7 +192,7 @@ class LossScaleOptimizer(optimizer.Optimizer):
         mapping = list(zip(self.variables, optimizer_variables))
         with backend.StatelessScope(state_mapping=mapping) as scope:
             self.step_counter.assign(0)
-            self.dynamic_scale.assign(self.dynamic_scale / 2.0)
+            self.dynamic_scale.assign(ops.multiply(self.dynamic_scale, 0.5))
         new_optimizer_variables = []
         for v in self.variables:
             new_optimizer_variables.append(scope.get_current_value(v))
@@ -190,7 +226,7 @@ class LossScaleOptimizer(optimizer.Optimizer):
 
         def upscale():
             self.step_counter.assign(0)
-            self.dynamic_scale.assign(self.dynamic_scale * 2.0)
+            self.dynamic_scale.assign(ops.multiply(self.dynamic_scale, 2.0))
 
         def increment():
             self.step_counter.assign_add(1)
@@ -205,7 +241,7 @@ class LossScaleOptimizer(optimizer.Optimizer):
     def _stateful_handle_non_finite_grads(self):
         # If any inf or nan in grads, downscale loss and reset counter.
         self.step_counter.assign(0)
-        self.dynamic_scale.assign(self.dynamic_scale / 2.0)
+        self.dynamic_scale.assign(ops.multiply(self.dynamic_scale, 0.5))
 
     def _common_apply(self, grads, trainable_variables=None):
         finite = self.check_finite(grads)
@@ -278,25 +314,22 @@ class LossScaleOptimizer(optimizer.Optimizer):
 
     def scale_loss(self, loss):
         scale = self.dynamic_scale if self.built else self.initial_scale
-        return loss * scale
+        return ops.multiply(loss, scale)
 
     def finalize_variable_values(self, var_list):
         self.inner_optimizer.finalize_variable_values(var_list)
 
     def get_config(self):
-        config = super().get_config()
+        # Do not use super().get_config() as only "name" is supported.
         inner_optimizer_config = serialization_lib.serialize_keras_object(
             self.inner_optimizer
         )
-        config.update(
-            {
-                "inner_optimizer": inner_optimizer_config,
-                "initial_scale": self.initial_scale,
-                "dynamic_growth_steps": self.dynamic_growth_steps,
-            }
-        )
-        del config["learning_rate"]
-        return config
+        return {
+            "name": self.name,
+            "inner_optimizer": inner_optimizer_config,
+            "initial_scale": self.initial_scale,
+            "dynamic_growth_steps": self.dynamic_growth_steps,
+        }
 
     @classmethod
     def from_config(cls, config, custom_objects=None):

keras/src/optimizers/muon.py

@@ -20,7 +20,7 @@ class Muon(optimizer.Optimizer):
     The Muon optimizer can use both the Muon update step or the
     AdamW update step based on the following:
 
-    - For any variable that isn't 2D, 3D or 4D, the AdamW step
+    - For any variable that isn't 2D, the AdamW step
         will be used. This is not configurable.
     - If the argument `exclude_embeddings` (defaults to `True`) is set
     to `True`, the AdamW step will be used.
@@ -46,10 +46,12 @@ class Muon(optimizer.Optimizer):
             that takes no arguments and returns the actual value to use.
             The exponential decay rate for the 1st moment estimates. Defaults to
             `0.9`.
-        adam_beta_2: A float value or a constant float tensor, ora callable
+        adam_beta_2: A float value or a constant float tensor, or a callable
             that takes no arguments and returns the actual value to use.
             The exponential decay rate for the 2nd moment estimates. Defaults to
             `0.999`.
+        adam_weight_decay: Float. If set, weight decay is applied when using
+            the Adam optimizer.
         epsilon: A small constant for numerical stability. This is
             "epsilon hat" in the Kingma and Ba paper
             (in the formula just before Section 2.1),
@@ -67,11 +69,15 @@ class Muon(optimizer.Optimizer):
             It is recommended to use the default value
         adam_lr_ratio: Float, the ratio of the learning rate when
                 using Adam to the main learning rate.
-                it is recommended to set it to 0.1
+                It is recommended to set it to 1
         momentum: Float, momentum used by internal SGD.
         ns_steps: Integer, number of Newton-Schulz iterations to run.
         nesterov: Boolean, whether to use Nesterov-style momentum
         {{base_optimizer_keyword_args}}
+        rms_rate: Float. A parameter from https://arxiv.org/abs/2502.16982
+            that can enhance the stability of Muon, allowing it to use the
+            same learning rate and weight decay as Adam. Defaults to `0.2`.
+            Set to `None` to disable this feature.
     """
 
     def __init__(
@@ -79,8 +85,9 @@ class Muon(optimizer.Optimizer):
         learning_rate=0.001,
         adam_beta_1=0.9,
         adam_beta_2=0.999,
+        adam_weight_decay=0.004,
         epsilon=1e-7,
-        weight_decay=0.1,
+        weight_decay=0.004,
         clipnorm=None,
         clipvalue=None,
         global_clipnorm=None,
@@ -95,10 +102,11 @@ class Muon(optimizer.Optimizer):
         muon_a=3.4445,
         muon_b=-4.7750,
         muon_c=2.0315,
-        adam_lr_ratio=0.1,
+        adam_lr_ratio=1,
         momentum=0.95,
-        ns_steps=6,
+        ns_steps=5,
         nesterov=True,
+        rms_rate=0.2,
         **kwargs,
     ):
         super().__init__(
@@ -127,12 +135,13 @@ class Muon(optimizer.Optimizer):
         self.nesterov = nesterov
         self.exclude_embeddings = exclude_embeddings
         self.exclude_layers = exclude_layers or []
+        self.adam_weight_decay = adam_weight_decay
+        self.rms_rate = rms_rate
 
     def _should_use_adamw(self, variable):
-        # To use it with 4D convolutional filters,
         # it works well to just flatten their last 3 dimensions.
         # any {0,1}-D parameters should all be optimized by adam
-        if not 1 < len(variable.shape) < 4:
+        if len(variable.shape) != 2:
             return True
         if self.exclude_embeddings and "embedding" in variable.path.lower():
             return True
@@ -153,52 +162,50 @@ class Muon(optimizer.Optimizer):
         if self.built:
             return
         super().build(var_list)
-        self.adam_momentums = {}
-        self.adam_velocities = {}
-
-        self.muon_momentums = {}
-        self.muon_velocities = {}
+        # Momentums are for both Muon and Adam
+        self.momentums = [None] * len(var_list)
+        # Velocities are just for Adam
+        self.adam_velocities = [None] * len(var_list)
 
         for var in var_list:
             if not self._overwrite_variable_with_gradient(var):
-                self.adam_momentums[var.path] = (
+                self.momentums[self._get_variable_index(var)] = (
                     self.add_variable_from_reference(
                         reference_variable=var, name="momentum"
                     )
                 )
                 if self._should_use_adamw(var):
-                    self.adam_velocities[var.path] = (
+                    self.adam_velocities[self._get_variable_index(var)] = (
                         self.add_variable_from_reference(
                             reference_variable=var, name="velocity"
                         )
                     )
 
     def update_step(self, gradient, variable, learning_rate):
-        if self._should_use_adamw(variable):
+        variable_index = self._get_variable_index(variable)
+        m = self.momentums[variable_index]
+        v = self.adam_velocities[variable_index]
+
+        # The presence of the velocity tells us that this variable is for Adam
+        if v is not None:
             # It should be noted that lr is one-tenth when using adamw.
             self._adamw_update_step(
-                gradient, variable, learning_rate * self.adam_lr_ratio
+                gradient, variable, learning_rate * self.adam_lr_ratio, m, v
             )
         else:
-            self._muon_update_step(gradient, variable, learning_rate)
+            self._muon_update_step(gradient, variable, learning_rate, m)
 
-    def _muon_update_step(self, gradient, variable, lr):
-        m = self.adam_momentums[variable.path]
+    def _muon_update_step(self, gradient, variable, lr, m):
         self.assign_add(m, ops.add(gradient, m * (self.momentum - 1)))
-        shape = variable.shape
         if self.nesterov:
             g = ops.add(gradient, self.momentum * m)
         else:
             g = m
+        update = self.zeropower_via_newtonschulz5(g, self.ns_steps)
 
-        self.assign_sub(
-            variable,
-            lr
-            * self.zeropower_via_newtonschulz5(g, self.ns_steps)
-            * max(1, shape[0] / shape[1]) ** 0.5,
-        )
+        self.assign_sub(variable, self.lr_adjust(lr * update))
 
-    def _adamw_update_step(self, gradient, variable, learning_rate):
+    def _adamw_update_step(self, gradient, variable, learning_rate, m, v):
         """Update step given gradient and the associated model variable."""
         lr = ops.cast(learning_rate, variable.dtype)
         gradient = ops.cast(gradient, variable.dtype)
@@ -210,9 +217,6 @@ class Muon(optimizer.Optimizer):
             ops.cast(self.adam_beta_2, variable.dtype), local_step
         )
 
-        m = self.adam_momentums[variable.path]
-        v = self.adam_velocities[variable.path]
-
         alpha = lr * ops.sqrt(1 - adam_beta_2_power) / (1 - adam_beta_1_power)
 
         self.assign_add(
@@ -239,6 +243,20 @@ class Muon(optimizer.Optimizer):
         X = ops.transpose(X, temp_order)
         return X
 
+    def lr_adjust(self, x):
+        """Adjusts learning rate based on the Moonlight implementation.
+        This method enhances the stability of Muon, allowing it to use the same
+        learning rate and weight decay as Adam. For details, see
+        https://arxiv.org/abs/2502.16982.
+        For a 2D matrix, the update is scaled by `sqrt(max(n, m)) * rms_rate`,
+        where `n` and `m` are the dimensions of the matrix.
+        """
+        if self.rms_rate is None:
+            return x
+        # moonlight version
+        # https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py
+        return x * ops.sqrt(ops.maximum(x.shape[0], x.shape[1])) * self.rms_rate
+
     def zeropower_via_newtonschulz5(self, x, steps: int):
         """We apply the Newton-Schulz iteration to compute matrix G.
 
@@ -268,6 +286,20 @@ class Muon(optimizer.Optimizer):
             x = self.transpose_last_axis(x)
         return x
 
+    def _apply_weight_decay(self, variables):
+        for variable in variables:
+            if not self._use_weight_decay(variable):
+                continue
+            if self._should_use_adamw(variable):
+                weight_decay_value = self.adam_weight_decay
+            else:
+                weight_decay_value = self.weight_decay
+            if weight_decay_value is None:
+                continue
+            wd = ops.cast(weight_decay_value, variable.dtype)
+            lr = ops.cast(self.learning_rate, variable.dtype)
+            variable.assign(variable - variable * wd * lr)
+
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -284,6 +316,8 @@ class Muon(optimizer.Optimizer):
                 "ns_steps": self.ns_steps,
                 "nesterov": self.nesterov,
                 "exclude_embeddings": self.exclude_embeddings,
+                "adam_weight_decay": self.adam_weight_decay,
+                "rms_rate": self.rms_rate,
             }
         )
         return config

keras/src/optimizers/schedules/learning_rate_schedule.py

@@ -584,9 +584,10 @@ class CosineDecay(LearningRateSchedule):
     schedule applies a linear increase per optimizer step to our learning rate
     from `initial_learning_rate` to `warmup_target` for a duration of
     `warmup_steps`. Afterwards, it applies a cosine decay function taking our
-    learning rate from `warmup_target` to `alpha` for a duration of
-    `decay_steps`. If `warmup_target` is None we skip warmup and our decay
-    will take our learning rate from `initial_learning_rate` to `alpha`.
+    learning rate from `warmup_target` to `warmup_target * alpha` for a
+    duration of `decay_steps`. If `warmup_target` is None we skip warmup and
+    our decay will take our learning rate from `initial_learning_rate` to
+    `initial_learning_rate * alpha`.
     It requires a `step` value to  compute the learning rate. You can
     just pass a backend variable that you increment at each training step.
 

keras/src/quantizers/__init__.py

@@ -1,6 +1,10 @@
 import inspect
 
 from keras.src.api_export import keras_export
+from keras.src.quantizers.quantization_config import Float8QuantizationConfig
+from keras.src.quantizers.quantization_config import Int4QuantizationConfig
+from keras.src.quantizers.quantization_config import Int8QuantizationConfig
+from keras.src.quantizers.quantization_config import QuantizationConfig
 from keras.src.quantizers.quantizers import AbsMaxQuantizer
 from keras.src.quantizers.quantizers import Quantizer
 from keras.src.quantizers.quantizers import abs_max_quantize
@@ -13,7 +17,14 @@ from keras.src.quantizers.quantizers import unpack_int4
 from keras.src.saving import serialization_lib
 from keras.src.utils.naming import to_snake_case
 
-ALL_OBJECTS = {Quantizer, AbsMaxQuantizer}
+ALL_OBJECTS = {
+    Quantizer,
+    AbsMaxQuantizer,
+    QuantizationConfig,
+    Int8QuantizationConfig,
+    Int4QuantizationConfig,
+    Float8QuantizationConfig,
+}
 ALL_OBJECTS_DICT = {cls.__name__: cls for cls in ALL_OBJECTS}
 ALL_OBJECTS_DICT.update(
     {to_snake_case(cls.__name__): cls for cls in ALL_OBJECTS}

keras/src/quantizers/gptq.py

@@ -2,6 +2,7 @@ import types
 from functools import partial
 
 from keras.src import ops
+from keras.src import quantizers
 from keras.src.layers import Dense
 from keras.src.layers import EinsumDense
 from keras.src.ops import linalg
@@ -295,12 +296,7 @@ class GPTQ:
             # For EinsumDense, we determine the effective 2D dimensions.
             self.kernel_shape = layer.kernel.shape
             shape = list(self.kernel_shape)
-            try:
-                d_model_dim_index = shape.index(max(shape))
-            except ValueError:
-                raise TypeError(
-                    f"Could not determine hidden dimension from shape {shape}"
-                )
+            d_model_dim_index = shape.index(max(shape))
 
             if d_model_dim_index == 0:  # QKV projection case
                 in_features, heads, head_dim = shape
@@ -476,6 +472,12 @@ class GPTQ:
             quantized, self.original_layer.quantized_kernel.dtype
         )
 
+        if self.config.weight_bits == 4:
+            # For 4-bit weights, we need to pack them into bytes
+            quantized, _, _ = quantizers.pack_int4(
+                quantized, axis=0, dtype="uint8"
+            )
+
         del self.original_layer._kernel
         self.original_layer.quantized_kernel.assign(quantized)
         self.original_layer.kernel_scale.assign(scale)

keras/src/quantizers/gptq_config.py

@@ -1,8 +1,9 @@
 from keras.src.api_export import keras_export
+from keras.src.quantizers.quantization_config import QuantizationConfig
 
 
 @keras_export("keras.quantizers.GPTQConfig")
-class GPTQConfig:
+class GPTQConfig(QuantizationConfig):
     """Configuration class for the GPTQ (Gradient-based Post-Training
     Quantization) algorithm.
 
@@ -131,6 +132,12 @@ class GPTQConfig:
         activation_order: (bool, optional) If `True`, reorders weight columns
             based on activation magnitude, which can improve quantization
             accuracy. Defaults to `False`.
+        quantization_layer_structure: (dict, optional) A dictionary defining the
+            model's quantization structure. It should contain:
+            - "pre_block_layers": list of layers to run before the first block.
+            - "sequential_blocks": list of blocks to be quantized sequentially.
+            If not provided, the model must implement
+            `get_quantization_layer_structure`.
     """
 
     def __init__(
@@ -146,7 +153,9 @@ class GPTQConfig:
         group_size: int = 128,
         symmetric: bool = False,
         activation_order: bool = False,
+        quantization_layer_structure: dict = None,
     ):
+        super().__init__()
         if weight_bits not in [2, 3, 4, 8]:
             raise ValueError(
                 f"Unsupported weight_bits {weight_bits}. "
@@ -174,6 +183,32 @@ class GPTQConfig:
         self.group_size = group_size
         self.symmetric = symmetric
         self.activation_order = activation_order
+        self.quantization_layer_structure = quantization_layer_structure
+
+    def get_config(self):
+        return {
+            # Dataset and Tokenizer are only required for a one-time
+            # calibration and are not saved in the config.
+            "dataset": None,
+            "tokenizer": None,
+            "weight_bits": self.weight_bits,
+            "num_samples": self.num_samples,
+            "per_channel": self.per_channel,
+            "sequence_length": self.sequence_length,
+            "hessian_damping": self.hessian_damping,
+            "group_size": self.group_size,
+            "symmetric": self.symmetric,
+            "activation_order": self.activation_order,
+            "quantization_layer_structure": self.quantization_layer_structure,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+    @property
+    def mode(self):
+        return "gptq"
 
     def dtype_policy_string(self):
         """Returns the dtype policy string for this configuration.