keras-nightly 3.14.0.dev2026011904__tar.gz → 3.14.0.dev2026012104__tar.gz

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/PKG-INFO

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

{keras_nightly-3.14.0.dev2026011904/keras/api → keras_nightly-3.14.0.dev2026012104/keras/api/_tf_keras/keras}/ops/__init__.py

@@ -245,6 +245,7 @@ from keras.src.ops.numpy import mod as mod
 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 nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
 from keras.src.ops.numpy import nextafter as nextafter

{keras_nightly-3.14.0.dev2026011904/keras/api → keras_nightly-3.14.0.dev2026012104/keras/api/_tf_keras/keras}/ops/numpy/__init__.py

@@ -129,6 +129,7 @@ from keras.src.ops.numpy import mod as mod
 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 nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
 from keras.src.ops.numpy import nextafter as nextafter

{keras_nightly-3.14.0.dev2026011904/keras/api/_tf_keras/keras → keras_nightly-3.14.0.dev2026012104/keras/api}/ops/__init__.py

@@ -245,6 +245,7 @@ from keras.src.ops.numpy import mod as mod
 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 nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
 from keras.src.ops.numpy import nextafter as nextafter

{keras_nightly-3.14.0.dev2026011904/keras/api/_tf_keras/keras → keras_nightly-3.14.0.dev2026012104/keras/api}/ops/numpy/__init__.py

@@ -129,6 +129,7 @@ from keras.src.ops.numpy import mod as mod
 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 nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
 from keras.src.ops.numpy import nextafter as nextafter

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/backend/jax/numpy.py

@@ -1013,6 +1013,11 @@ def moveaxis(x, source, destination):
     return jnp.moveaxis(x, source=source, destination=destination)
 
 
+def nansum(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    return jnp.nansum(x, axis=axis, keepdims=keepdims)
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     x = convert_to_tensor(x)
     return jnp.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/backend/numpy/numpy.py

@@ -960,6 +960,17 @@ def moveaxis(x, source, destination):
     return np.moveaxis(x, source=source, destination=destination)
 
 
+def nansum(x, axis=None, keepdims=False):
+    axis = standardize_axis_for_numpy(axis)
+    dtype = standardize_dtype(x.dtype)
+
+    if dtype in ("bool", "int8", "int16"):
+        dtype = "int32"
+    elif dtype in ("uint8", "uint16"):
+        dtype = "uint32"
+    return np.nansum(x, axis=axis, keepdims=keepdims).astype(dtype)
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     return np.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)
 

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/backend/openvino/numpy.py

@@ -2056,6 +2056,10 @@ def moveaxis(x, source, destination):
     return OpenVINOKerasTensor(ov_opset.transpose(x, axes_const).output(0))
 
 
+def nansum(x, axis=None, keepdims=False):
+    raise NotImplementedError("`nansum` is not supported with openvino backend")
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     x = get_ov_output(x)
     dtype = x.get_element_type()

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/backend/tensorflow/numpy.py

@@ -2125,6 +2125,22 @@ def moveaxis(x, source, destination):
     return tf.transpose(x, perm)
 
 
+def nansum(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    dtype = standardize_dtype(x.dtype)
+    x_clean = tf.where(
+        tf.math.is_nan(cast(x, config.floatx())), tf.zeros((), dtype=dtype), x
+    )
+
+    if dtype in ("bool", "int8", "int16"):
+        dtype = "int32"
+    elif dtype in ("uint8", "uint16"):
+        dtype = "uint32"
+    x_clean = cast(x_clean, dtype)
+
+    return tf.reduce_sum(x_clean, axis=axis, keepdims=keepdims)
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     x = convert_to_tensor(x)
 

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/backend/torch/numpy.py

@@ -1272,6 +1272,20 @@ def moveaxis(x, source, destination):
     return torch.moveaxis(x, source=source, destination=destination)
 
 
+def nansum(x, axis=None, keepdims=False):
+    if isinstance(x, (list, tuple)):
+        x = stack(x)
+    x = convert_to_tensor(x)
+    dtype = standardize_dtype(x.dtype)
+
+    if dtype in ("bool", "uint8", "int8", "int16"):
+        dtype = "int32"
+
+    if axis == () or axis == []:
+        return cast(torch.nan_to_num(x, nan=0), dtype)
+    return cast(torch.nansum(x, dim=axis, keepdim=keepdims), dtype)
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     x = convert_to_tensor(x)
     return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/layers/preprocessing/image_preprocessing/center_crop.py

@@ -183,28 +183,26 @@ class CenterCrop(BaseImagePreprocessingLayer):
 
     def transform_images(self, images, transformation=None, training=True):
         inputs = self.backend.cast(images, self.compute_dtype)
+        inputs_shape = self.backend.shape(inputs)
+
         if self.data_format == "channels_first":
-            init_height = inputs.shape[-2]
-            init_width = inputs.shape[-1]
+            init_height = inputs_shape[-2]
+            init_width = inputs_shape[-1]
         else:
-            init_height = inputs.shape[-3]
-            init_width = inputs.shape[-2]
-
-        if init_height is None or init_width is None:
-            # Dynamic size case. TODO.
-            raise ValueError(
-                "At this time, CenterCrop can only "
-                "process images with a static spatial "
-                f"shape. Received: inputs.shape={inputs.shape}"
-            )
+            init_height = inputs_shape[-3]
+            init_width = inputs_shape[-2]
 
+        # All these operations work both with ints (static sizes) and scalar
+        # tensors (dynamic sizes).
         h_diff = init_height - self.height
         w_diff = init_width - self.width
 
-        h_start = int(h_diff / 2)
-        w_start = int(w_diff / 2)
+        h_start = h_diff // 2
+        w_start = w_diff // 2
 
-        if h_diff >= 0 and w_diff >= 0:
+        if (not isinstance(h_diff, int) or h_diff >= 0) and (
+            not isinstance(w_diff, int) or w_diff >= 0
+        ):
             if len(inputs.shape) == 4:
                 if self.data_format == "channels_first":
                     return inputs[

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/ops/image.py

@@ -565,6 +565,7 @@ class ExtractPatches(Operation):
         if isinstance(size, int):
             size = (size, size)
         self.size = size
+        self.is_3d = len(self.size) == 3
         if strides is None:
             strides = size
         self.strides = strides
@@ -588,11 +589,21 @@ class ExtractPatches(Operation):
         if self.data_format == "channels_last":
             channels_in = images_shape[-1]
         else:
-            channels_in = images_shape[-3]
-        if original_ndim == 3:
-            images_shape = [1] + images_shape
-        filters = self.size[0] * self.size[1] * channels_in
-        kernel_size = (self.size[0], self.size[1])
+            channels_in = images_shape[-4] if self.is_3d else images_shape[-3]
+
+        if self.is_3d:
+            # 3D patch extraction
+            if original_ndim == 4:
+                images_shape = [1] + images_shape
+            filters = self.size[0] * self.size[1] * self.size[2] * channels_in
+            kernel_size = (self.size[0], self.size[1], self.size[2])
+        else:
+            # 2D patch extraction
+            if original_ndim == 3:
+                images_shape = [1] + images_shape
+            filters = self.size[0] * self.size[1] * channels_in
+            kernel_size = (self.size[0], self.size[1])
+
         out_shape = compute_conv_output_shape(
             images_shape,
             filters,
@@ -602,10 +613,24 @@ class ExtractPatches(Operation):
             data_format=self.data_format,
             dilation_rate=self.dilation_rate,
         )
-        if original_ndim == 3:
-            out_shape = out_shape[1:]
+
+        if self.is_3d:
+            if original_ndim == 4:
+                out_shape = out_shape[1:]
+        else:
+            if original_ndim == 3:
+                out_shape = out_shape[1:]
         return KerasTensor(shape=out_shape, dtype=images.dtype)
 
+    def get_config(self):
+        return {
+            "size": self.size,
+            "strides": self.strides,
+            "dilation_rate": self.dilation_rate,
+            "padding": self.padding,
+            "data_format": self.data_format,
+        }
+
 
 @keras_export("keras.ops.image.extract_patches")
 def extract_patches(
@@ -687,26 +712,6 @@ def extract_patches(
                 f"Received: size={size} with length {len(size)}"
             )
 
-    # Determine 2D vs 3D based on size argument
-    if not isinstance(size, int) and len(size) == 3:
-        # 3D patch extraction
-        if any_symbolic_tensors((images,)):
-            return ExtractPatches3D(
-                size=size,
-                strides=strides,
-                dilation_rate=dilation_rate,
-                padding=padding,
-                data_format=data_format,
-            ).symbolic_call(images)
-        return _extract_patches_3d(
-            images,
-            size,
-            strides,
-            dilation_rate,
-            padding,
-            data_format=data_format,
-        )
-
     # 2D patch extraction (default)
     if any_symbolic_tensors((images,)):
         return ExtractPatches(
@@ -729,6 +734,23 @@ def _extract_patches(
     dilation_rate=1,
     padding="valid",
     data_format=None,
+):
+    if not isinstance(size, int) and len(size) == 3:
+        return _extract_patches_3d(
+            images, size, strides, dilation_rate, padding, data_format
+        )
+    return _extract_patches_2d(
+        images, size, strides, dilation_rate, padding, data_format
+    )
+
+
+def _extract_patches_2d(
+    images,
+    size,
+    strides=None,
+    dilation_rate=1,
+    padding="valid",
+    data_format=None,
 ):
     if isinstance(size, int):
         patch_h = patch_w = size
@@ -768,74 +790,6 @@ def _extract_patches(
     return patches
 
 
-class ExtractPatches3D(Operation):
-    def __init__(
-        self,
-        size,
-        strides=None,
-        dilation_rate=1,
-        padding="valid",
-        data_format=None,
-        *,
-        name=None,
-    ):
-        super().__init__(name=name)
-        if isinstance(size, int):
-            size = (size, size, size)
-        elif len(size) != 3:
-            raise TypeError(
-                "Invalid `size` argument. Expected an "
-                f"int or a tuple of length 3. Received: size={size}"
-            )
-        self.size = size
-        if strides is not None:
-            if isinstance(strides, int):
-                strides = (strides, strides, strides)
-            elif len(strides) != 3:
-                raise ValueError(f"Invalid `strides` argument. Got: {strides}")
-        else:
-            strides = size
-        self.strides = strides
-        self.dilation_rate = dilation_rate
-        self.padding = padding
-        self.data_format = backend.standardize_data_format(data_format)
-
-    def call(self, volumes):
-        return _extract_patches_3d(
-            volumes,
-            self.size,
-            self.strides,
-            self.dilation_rate,
-            self.padding,
-            self.data_format,
-        )
-
-    def compute_output_spec(self, volumes):
-        volumes_shape = list(volumes.shape)
-        original_ndim = len(volumes_shape)
-        strides = self.strides
-        if self.data_format == "channels_last":
-            channels_in = volumes_shape[-1]
-        else:
-            channels_in = volumes_shape[-4]
-        if original_ndim == 4:
-            volumes_shape = [1] + volumes_shape
-        filters = self.size[0] * self.size[1] * self.size[2] * channels_in
-        kernel_size = (self.size[0], self.size[1], self.size[2])
-        out_shape = compute_conv_output_shape(
-            volumes_shape,
-            filters,
-            kernel_size,
-            strides=strides,
-            padding=self.padding,
-            data_format=self.data_format,
-            dilation_rate=self.dilation_rate,
-        )
-        if original_ndim == 4:
-            out_shape = out_shape[1:]
-        return KerasTensor(shape=out_shape, dtype=volumes.dtype)
-
-
 def _extract_patches_3d(
     volumes,
     size,
@@ -935,8 +889,11 @@ def extract_patches_3d(
     >>> patches.shape
     (3, 3, 3, 81)
     """
+    # Convert int to 3-tuple for 3D
+    if isinstance(size, int):
+        size = (size, size, size)
     if any_symbolic_tensors((volumes,)):
-        return ExtractPatches3D(
+        return ExtractPatches(
             size=size,
             strides=strides,
             dilation_rate=dilation_rate,

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/ops/numpy.py

@@ -5064,6 +5064,67 @@ def moveaxis(x, source, destination):
     return backend.numpy.moveaxis(x, source=source, destination=destination)
 
 
+class Nansum(Operation):
+    def __init__(self, axis=None, keepdims=False, *, name=None):
+        super().__init__(name=name)
+        self.axis = axis
+        self.keepdims = keepdims
+
+    def call(self, x):
+        return backend.numpy.nansum(x, axis=self.axis, keepdims=self.keepdims)
+
+    def compute_output_spec(self, x):
+        dtype = dtypes.result_type(getattr(x, "dtype", backend.floatx()))
+
+        if dtype in ("bool", "int8", "int16"):
+            dtype = "int32"
+        elif dtype in ("uint8", "uint16"):
+            dtype = "uint32"
+
+        if backend.backend() == "torch" and dtype == "uint32":
+            dtype = "int32"
+        sparse = getattr(x, "sparse", False)
+        return KerasTensor(
+            reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
+            dtype=dtype,
+            sparse=sparse,
+        )
+
+
+@keras_export(["keras.ops.nansum", "keras.ops.numpy.nansum"])
+def nansum(x, axis=None, keepdims=False):
+    """Sum of a tensor over the given axes, ignoring NaNs.
+
+    Args:
+        x: Input tensor.
+        axis: Axis or axes along which the sum is computed. The default is to
+            compute the sum of the flattened tensor.
+        keepdims: If this is set to `True`, the axes which are reduced are left
+            in the result as dimensions with size one.
+
+    Returns:
+        Output tensor containing the sum, with NaN values ignored.
+
+    Examples:
+    >>> import numpy as np
+    >>> from keras import ops
+    >>> x = np.array([[1.0, np.nan, 3.0],
+    ...               [np.nan, 2.0, 1.0]])
+    >>> ops.nansum(x)
+    7.0
+
+    >>> ops.nansum(x, axis=1)
+    array([4., 3.])
+
+    >>> ops.nansum(x, axis=1, keepdims=True)
+    array([[4.],
+           [3.]])
+    """
+    if any_symbolic_tensors((x,)):
+        return Nansum(axis=axis, keepdims=keepdims).symbolic_call(x)
+    return backend.numpy.nansum(x, axis=axis, keepdims=keepdims)
+
+
 class NanToNum(Operation):
     def __init__(self, nan=0.0, posinf=None, neginf=None, *, name=None):
         super().__init__(name=name)

keras_nightly-3.14.0.dev2026012104/keras/src/saving/orbax_util.py

@@ -0,0 +1,50 @@
+"""Orbax checkpoint loading functionality."""
+
+import os
+
+from keras.src.utils.module_utils import ocp
+
+
+def is_orbax_checkpoint(filepath):
+    """Check if the given path is an Orbax checkpoint directory.
+
+    This function implements custom detection logic instead of relying on
+    Orbax APIs which may be unreliable in some environments.
+    """
+    if not os.path.exists(filepath) or not os.path.isdir(filepath):
+        return False
+
+    try:
+        # List directory contents
+        contents = os.listdir(filepath)
+
+        # A set is more efficient for membership testing
+        orbax_indicators = {
+            "orbax.checkpoint",
+            "pytree.orbax-checkpoint",
+            "checkpoint_metadata",
+        }
+
+        # Fast check for standard files
+        if not orbax_indicators.isdisjoint(contents):
+            return True
+
+        # Check for step directories or temporary files in a single pass
+        return any(
+            ".orbax-checkpoint-tmp" in item
+            or (item.isdigit() and os.path.isdir(os.path.join(filepath, item)))
+            for item in contents
+        )
+
+    except (OSError, PermissionError):
+        # If we can't read the directory, assume it's not a checkpoint
+        return False
+
+
+def find_latest_orbax_checkpoint(checkpoint_dir):
+    """Find the latest checkpoint in an Orbax checkpoint directory."""
+    checkpointer = ocp.training.Checkpointer(directory=checkpoint_dir)
+    latest = checkpointer.latest
+    if latest is None:
+        raise ValueError(f"No valid checkpoints found in {checkpoint_dir}")
+    return os.path.join(checkpoint_dir, str(latest.step))

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras/src/version.py

@@ -1,7 +1,7 @@
 from keras.src.api_export import keras_export
 
 # Unique source of truth for the version number.
-__version__ = "3.14.0.dev2026011904"
+__version__ = "3.14.0.dev2026012104"
 
 
 @keras_export("keras.version")

{keras_nightly-3.14.0.dev2026011904 → keras_nightly-3.14.0.dev2026012104}/keras_nightly.egg-info/PKG-INFO

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

keras_nightly-3.14.0.dev2026011904/keras/src/saving/orbax_util.py

@@ -1,26 +0,0 @@
-"""Orbax checkpoint loading functionality."""
-
-import os
-
-from keras.src.utils.module_utils import ocp
-
-
-def is_orbax_checkpoint(filepath):
-    """Check if the given path is an Orbax checkpoint directory."""
-    if not os.path.exists(filepath):
-        return False
-
-    try:
-        return ocp.is_orbax_checkpoint(filepath)
-    except (ImportError, AttributeError):
-        # Fallback to check for orbax.checkpoint file if Orbax API not available
-        return os.path.isfile(os.path.join(filepath, "orbax.checkpoint"))
-
-
-def find_latest_orbax_checkpoint(checkpoint_dir):
-    """Find the latest checkpoint in an Orbax checkpoint directory."""
-    checkpointer = ocp.training.Checkpointer(directory=checkpoint_dir)
-    latest = checkpointer.latest
-    if latest is None:
-        raise ValueError(f"No valid checkpoints found in {checkpoint_dir}")
-    return os.path.join(checkpoint_dir, str(latest.step))