keras-nightly 3.14.0.dev2026012704__tar.gz → 3.14.0.dev2026012904__tar.gz

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-nightly
-Version: 3.14.0.dev2026012704
+Version: 3.14.0.dev2026012904
 Summary: Multi-backend Keras
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/api/_tf_keras/keras/dtype_policies/__init__.py

@@ -17,6 +17,9 @@ from keras.src.dtype_policies.dtype_policy import (
 from keras.src.dtype_policies.dtype_policy import (
     GPTQDTypePolicy as GPTQDTypePolicy,
 )
+from keras.src.dtype_policies.dtype_policy import (
+    Int4DTypePolicy as Int4DTypePolicy,
+)
 from keras.src.dtype_policies.dtype_policy import (
     QuantizedDTypePolicy as QuantizedDTypePolicy,
 )

{keras_nightly-3.14.0.dev2026012704/keras/api → keras_nightly-3.14.0.dev2026012904/keras/api/_tf_keras/keras}/ops/__init__.py

@@ -246,6 +246,7 @@ from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
 from keras.src.ops.numpy import nanmax as nanmax
+from keras.src.ops.numpy import nanmean as nanmean
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim

{keras_nightly-3.14.0.dev2026012704/keras/api → keras_nightly-3.14.0.dev2026012904/keras/api/_tf_keras/keras}/ops/numpy/__init__.py

@@ -130,6 +130,7 @@ from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
 from keras.src.ops.numpy import nanmax as nanmax
+from keras.src.ops.numpy import nanmean as nanmean
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/api/_tf_keras/keras/quantizers/__init__.py

@@ -24,6 +24,9 @@ from keras.src.quantizers.quantization_config import (
 from keras.src.quantizers.quantizers import AbsMaxQuantizer as AbsMaxQuantizer
 from keras.src.quantizers.quantizers import Quantizer as Quantizer
 from keras.src.quantizers.quantizers import abs_max_quantize as abs_max_quantize
+from keras.src.quantizers.quantizers import (
+    abs_max_quantize_grouped_with_zero_point as abs_max_quantize_grouped_with_zero_point,
+)
 from keras.src.quantizers.quantizers import (
     compute_float8_amax_history as compute_float8_amax_history,
 )

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/api/dtype_policies/__init__.py

@@ -17,6 +17,9 @@ from keras.src.dtype_policies.dtype_policy import (
 from keras.src.dtype_policies.dtype_policy import (
     GPTQDTypePolicy as GPTQDTypePolicy,
 )
+from keras.src.dtype_policies.dtype_policy import (
+    Int4DTypePolicy as Int4DTypePolicy,
+)
 from keras.src.dtype_policies.dtype_policy import (
     QuantizedDTypePolicy as QuantizedDTypePolicy,
 )

{keras_nightly-3.14.0.dev2026012704/keras/api/_tf_keras/keras → keras_nightly-3.14.0.dev2026012904/keras/api}/ops/__init__.py

@@ -246,6 +246,7 @@ from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
 from keras.src.ops.numpy import nanmax as nanmax
+from keras.src.ops.numpy import nanmean as nanmean
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim

{keras_nightly-3.14.0.dev2026012704/keras/api/_tf_keras/keras → keras_nightly-3.14.0.dev2026012904/keras/api}/ops/numpy/__init__.py

@@ -130,6 +130,7 @@ from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
 from keras.src.ops.numpy import nanmax as nanmax
+from keras.src.ops.numpy import nanmean as nanmean
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/api/quantizers/__init__.py

@@ -24,6 +24,9 @@ from keras.src.quantizers.quantization_config import (
 from keras.src.quantizers.quantizers import AbsMaxQuantizer as AbsMaxQuantizer
 from keras.src.quantizers.quantizers import Quantizer as Quantizer
 from keras.src.quantizers.quantizers import abs_max_quantize as abs_max_quantize
+from keras.src.quantizers.quantizers import (
+    abs_max_quantize_grouped_with_zero_point as abs_max_quantize_grouped_with_zero_point,
+)
 from keras.src.quantizers.quantizers import (
     compute_float8_amax_history as compute_float8_amax_history,
 )

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/backend/jax/core.py

@@ -98,7 +98,7 @@ if config.is_nnx_enabled():
         ):
             # Ensure 'mutable' is in nnx_metadata, but explicit 'mutable'
             # param takes precedence.
-            nnx_metadata["mutable"] = trainable if mutable is None else mutable
+            nnx_metadata["mutable"] = True if mutable is None else mutable
 
             # First, initialize a basic nnx.Variable with a dummy value
             # This sets up the NNX variable structure
@@ -603,7 +603,17 @@ def random_seed_dtype():
 
 
 def custom_gradient(fun):
-    return jax.custom_gradient(fun=fun)
+    fun_with_custom_gradient = jax.custom_gradient(fun=fun)
+
+    # Add a wrapper to unwrap variables, otherwise custom_gradient will fail
+    def fun_with_custom_gradient_wrapper(*args, **kwargs):
+        args, kwargs = tree.map_shape_structure(
+            lambda x: x.value if isinstance(x, KerasVariable) else x,
+            (args, kwargs),
+        )
+        return fun_with_custom_gradient(*args, **kwargs)
+
+    return fun_with_custom_gradient_wrapper
 
 
 def remat(f):

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/backend/jax/numpy.py

@@ -1018,6 +1018,11 @@ def nanmax(x, axis=None, keepdims=False):
     return jnp.nanmax(x, axis=axis, keepdims=keepdims)
 
 
+def nanmean(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    return jnp.nanmean(x, axis=axis, keepdims=keepdims)
+
+
 def nanmin(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     return jnp.nanmin(x, axis=axis, keepdims=keepdims)

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/backend/numpy/numpy.py

@@ -964,6 +964,11 @@ def nanmax(x, axis=None, keepdims=False):
     return np.nanmax(x, axis=axis, keepdims=keepdims)
 
 
+def nanmean(x, axis=None, keepdims=False):
+    dtype = dtypes.result_type(standardize_dtype(x.dtype), float)
+    return np.nanmean(x, axis=axis, keepdims=keepdims).astype(dtype)
+
+
 def nanmin(x, axis=None, keepdims=False):
     return np.nanmin(x, axis=axis, keepdims=keepdims)
 

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/backend/openvino/numpy.py

@@ -2111,6 +2111,12 @@ def nanmax(x, axis=None, keepdims=False):
     raise NotImplementedError("`nanmax` is not supported with openvino backend")
 
 
+def nanmean(x, axis=None, keepdims=False):
+    raise NotImplementedError(
+        "`nanmean` is not supported with openvino backend"
+    )
+
+
 def nanmin(x, axis=None, keepdims=False):
     raise NotImplementedError("`nanmin` is not supported with openvino backend")
 

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/backend/tensorflow/numpy.py

@@ -2145,6 +2145,27 @@ def nanmax(x, axis=None, keepdims=False):
     )
 
 
+def nanmean(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+
+    if axis == () or axis == []:
+        return x
+
+    if not x.dtype.is_floating:
+        return tf.reduce_mean(
+            tf.cast(x, "float32"), axis=axis, keepdims=keepdims
+        )
+
+    dtype = dtypes.result_type(standardize_dtype(x.dtype), float)
+    total_sum = cast(nansum(x, axis=axis, keepdims=keepdims), dtype)
+    normalizer = tf.reduce_sum(
+        cast(~tf.math.is_nan(x), dtype),
+        axis=axis,
+        keepdims=keepdims,
+    )
+    return tf.divide(total_sum, normalizer)
+
+
 def nanmin(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
 

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/backend/torch/numpy.py

@@ -1290,6 +1290,16 @@ def nanmax(x, axis=None, keepdims=False):
     )
 
 
+def nanmean(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+
+    if axis == () or axis == []:
+        return x
+
+    dtype = dtypes.result_type(standardize_dtype(x.dtype), float)
+    return torch.nanmean(cast(x, dtype), dim=axis, keepdim=keepdims)
+
+
 def nanmin(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     if not torch.is_floating_point(x):

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/callbacks/orbax_checkpoint.py

@@ -8,6 +8,7 @@ from keras.src.api_export import keras_export
 from keras.src.callbacks.monitor_callback import (
     MonitorCallback,  # For metric monitoring logic
 )
+from keras.src.saving import saving_lib
 from keras.src.utils.module_utils import ocp
 
 # Context and AsyncOptions are accessed through the lazy-loaded ocp module
@@ -117,6 +118,7 @@ class OrbaxCheckpoint(MonitorCallback):
         initial_value_threshold=None,
         max_to_keep=1,
         save_on_background=True,
+        save_weights_only=False,
     ):
         # Ensure orbax is available
         ocp.initialize()
@@ -131,10 +133,12 @@ class OrbaxCheckpoint(MonitorCallback):
         self.save_freq = save_freq
         self.max_to_keep = max_to_keep
         self.save_on_background = save_on_background
+        self.save_weights_only = save_weights_only
         self._batches_seen_since_last_saving = 0
         self._last_batch_seen = 0
         self._current_epoch = 0  # Keep track of epoch
         self._total_batches_seen = 0  # Global batch counter for step tracking
+        self._async_futures = []  # Track async save futures
 
         # Multi-host support
         self._multihost_initialized = self._is_multihost_initialized()
@@ -167,9 +171,14 @@ class OrbaxCheckpoint(MonitorCallback):
 
         # Create the V1 Checkpointer with direct parameter passing
         # Orbax will handle directory creation on all processes as needed
+        # save_decision_policy is required for proper coordination of
+        # rapid async saves
         self.checkpointer = ocp.training.Checkpointer(
             directory=directory,
             preservation_policy=preservation_policy,
+            save_decision_policy=ocp.training.save_decision_policies.FixedIntervalPolicy(
+                1
+            ),
         )
 
     def _is_multihost_initialized(self):
@@ -246,15 +255,35 @@ class OrbaxCheckpoint(MonitorCallback):
 
         # Save the nested state structures directly (preserving layer
         # names and structure)
-        composite_state = state_tree
+        if self.save_weights_only:
+            composite_state = {
+                "trainable_variables": state_tree["trainable_variables"],
+                "non_trainable_variables": state_tree[
+                    "non_trainable_variables"
+                ],
+            }
+        else:
+            composite_state = state_tree
+            # Include model configuration for full model restoration
+            # Use saving_lib helper to properly handle shared objects
+            config_json, _ = saving_lib._serialize_model_as_json(self.model)
+            composite_state["model_config"] = config_json
 
         # Use a single with statement. If context_options is empty,
         # Context() uses defaults.
         with ocp.Context():
-            if self.save_on_background:
-                self.checkpointer.save_pytree_async(step, composite_state)
-            else:
+            # Determine sync vs async based on save_on_background setting
+            use_sync = not self.save_on_background
+
+            if use_sync:
+                # Synchronous save
                 self.checkpointer.save_pytree(step, composite_state)
+            else:
+                # Async save
+                future = self.checkpointer.save_pytree_async(
+                    step, composite_state
+                )
+                self._async_futures.append(future)
 
     def on_train_batch_end(self, batch, logs=None):
         if self._should_save_on_batch(batch):
@@ -306,12 +335,11 @@ class OrbaxCheckpoint(MonitorCallback):
 
         if should_save:
             # Use epoch number as the step for Orbax save
-            # Keras has already made the save decision - Checkpointer will
-            # save unconditionally
             self._save_checkpoint(step=epoch, logs=logs)
 
     def on_train_end(self, logs=None):
-        # Close the Checkpointer to ensure all pending saves complete
+        # Close the Checkpointer - this waits for any pending async saves
+        # to complete before closing
         try:
             self.checkpointer.close()
         except Exception:
@@ -325,7 +353,12 @@ class OrbaxCheckpoint(MonitorCallback):
         This method blocks until all asynchronous checkpoint save operations
         have completed across all hosts in a multi-host setup.
         """
-        # Wait for any async operations to complete on this host
+        # Wait for all tracked async futures to complete
+        for future in self._async_futures:
+            future.result()  # Wait for completion
+        self._async_futures.clear()  # Clear completed futures
+
+        # Wait for any remaining async operations to complete on this host
         self.checkpointer.wait()
 
         # Multi-host synchronization: ensure all hosts complete

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/dtype_policies/__init__.py

@@ -6,6 +6,7 @@ from keras.src.dtype_policies.dtype_policy import AWQDTypePolicy
 from keras.src.dtype_policies.dtype_policy import DTypePolicy
 from keras.src.dtype_policies.dtype_policy import FloatDTypePolicy
 from keras.src.dtype_policies.dtype_policy import GPTQDTypePolicy
+from keras.src.dtype_policies.dtype_policy import Int4DTypePolicy
 from keras.src.dtype_policies.dtype_policy import QuantizedDTypePolicy
 from keras.src.dtype_policies.dtype_policy import QuantizedFloat8DTypePolicy
 from keras.src.dtype_policies.dtype_policy_map import DTypePolicyMap
@@ -18,6 +19,7 @@ ALL_OBJECTS = {
     QuantizedFloat8DTypePolicy,
     DTypePolicyMap,
     GPTQDTypePolicy,
+    Int4DTypePolicy,
 }
 ALL_OBJECTS_DICT = {cls.__name__: cls for cls in ALL_OBJECTS}
 

{keras_nightly-3.14.0.dev2026012704 → keras_nightly-3.14.0.dev2026012904}/keras/src/dtype_policies/dtype_policy.py

@@ -288,6 +288,79 @@ class QuantizedFloat8DTypePolicy(QuantizedDTypePolicy):
         return config
 
 
+@keras_export("keras.dtype_policies.Int4DTypePolicy")
+class Int4DTypePolicy(QuantizedDTypePolicy):
+    """Quantized dtype policy for int4 quantization.
+
+    This policy helps propagate quantization settings for int4 sub-channel
+    quantization when loading a quantized model in Keras format.
+
+    Args:
+        mode: The quantization mode. This should be a string in the format
+            `"int4/<block_size>"`.
+            -   `"int4"`: The identifier for the quantization algorithm.
+            -   `<block_size>`: The block size for sub-channel quantization.
+                Use -1 for per-channel (legacy) quantization. Any positive
+                integer enables sub-channel quantization with that block size.
+            Example: `"int4/128"` for sub-channel with 128-element groups.
+        source_name: The source dtype policy name, e.g. "float32".
+    """
+
+    def __init__(
+        self,
+        mode,
+        source_name=None,
+    ):
+        parts = mode.split("/")
+        expected_format = "'int4/<block_size>'"
+
+        # Validate format
+        if len(parts) != 2 or parts[0] != "int4":
+            raise ValueError(
+                "Invalid mode for Int4DTypePolicy. Expected format "
+                f"{expected_format}, but got '{mode}'."
+            )
+
+        # Validate and cast block_size
+        try:
+            block_size = int(parts[1])
+        except ValueError:
+            raise ValueError(
+                "Invalid mode for Int4DTypePolicy. <block_size> must be an "
+                f"integer. Expected format {expected_format}, but got '{mode}'."
+            )
+
+        # Validate supported values
+        if block_size < -1 or block_size == 0:
+            raise ValueError(
+                "Invalid block_size in mode. Supported values are "
+                "-1 (per-channel) or a positive integer (sub-channel), "
+                f"but got {block_size} from '{mode}'."
+            )
+
+        base_mode = parts[0]
+        super().__init__(
+            mode=base_mode,
+            source_name=source_name,
+        )
+
+        self._name = f"{mode}_from_{source_name}"
+        self.mode = base_mode
+        self.block_size = block_size
+
+    def __eq__(self, other):
+        if super().__eq__(other) is False:
+            return False
+        return self.block_size == other.block_size
+
+    def get_config(self):
+        config = super().get_config()
+        # Reconstruct the full mode string for serialization
+        mode = f"{self.mode}/{self.block_size}"
+        config.update({"mode": mode})
+        return config
+
+
 @keras_export("keras.dtype_policies.GPTQDTypePolicy")
 class GPTQDTypePolicy(QuantizedDTypePolicy):
     """Quantized dtype policy for GPTQ quantization.
@@ -525,8 +598,14 @@ def _get_quantized_dtype_policy_by_str(policy):
             f"Received: policy={policy}"
         )
     mode, source_name = split_name
-    if policy.startswith("int8") or policy.startswith("int4"):
+    if policy.startswith("int8"):
         return QuantizedDTypePolicy(mode, source_name)
+    elif policy.startswith("int4"):
+        # Check if mode has block_size component (e.g., "int4/128")
+        if "/" in mode:
+            return Int4DTypePolicy(mode, source_name)
+        else:
+            return QuantizedDTypePolicy(mode, source_name)
     elif policy.startswith("gptq"):
         return GPTQDTypePolicy(mode, source_name)
     elif policy.startswith("awq"):