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

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

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

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

@@ -2270,7 +2270,16 @@ def prod(x, axis=None, keepdims=False, dtype=None):
 
 
 def ptp(x, axis=None, keepdims=False):
-    raise NotImplementedError("`ptp` is not supported with openvino backend")
+    if axis == ():
+        return zeros_like(x)
+    x = get_ov_output(x)
+
+    x_resolved, resolved_axis = _resolve_axis(x, axis)
+
+    max_val = ov_opset.reduce_max(x_resolved, resolved_axis, keepdims)
+    min_val = ov_opset.reduce_min(x_resolved, resolved_axis, keepdims)
+
+    return OpenVINOKerasTensor(ov_opset.subtract(max_val, min_val).output(0))
 
 
 def quantile(x, q, axis=None, method="linear", keepdims=False):
@@ -2682,7 +2691,23 @@ def tensordot(x1, x2, axes=2):
 
 
 def round(x, decimals=0):
-    raise NotImplementedError("`round` is not supported with openvino backend")
+    x = get_ov_output(x)
+    x_type = x.get_element_type()
+    if x_type.is_integral() or x_type == Type.boolean:
+        x = ov_opset.convert(x, OPENVINO_DTYPES[config.floatx()])
+
+    if decimals == 0:
+        result = ov_opset.round(x, "half_to_even")
+    else:
+        factor = ov_opset.constant(10.0**decimals, x.get_element_type())
+        scaled = ov_opset.multiply(x, factor)
+        rounded = ov_opset.round(scaled, "half_to_even")
+        result = ov_opset.divide(rounded, factor)
+
+    if x_type.is_integral():
+        result = ov_opset.convert(result, x_type)
+
+    return OpenVINOKerasTensor(result.output(0))
 
 
 def tile(x, repeats):
@@ -2803,7 +2828,17 @@ def vdot(x1, x2):
 
 
 def vstack(xs):
-    raise NotImplementedError("`vstack` is not supported with openvino backend")
+    if not isinstance(xs, (list, tuple)):
+        xs = (xs,)
+    elems = [convert_to_tensor(elem) for elem in xs]
+    element_type = elems[0].output.get_element_type()
+    elems = [get_ov_output(elem, element_type) for elem in elems]
+    axis = 0
+    for i in range(1, len(elems)):
+        elems[0], elems[i] = _align_operand_types(
+            elems[0], elems[i], "vstack()"
+        )
+    return OpenVINOKerasTensor(ov_opset.concat(elems, axis).output(0))
 
 
 def vectorize(pyfunc, *, excluded=None, signature=None):
@@ -2859,9 +2894,20 @@ def divide(x1, x2):
 
 
 def divide_no_nan(x1, x2):
-    raise NotImplementedError(
-        "`divide_no_nan` is not supported with openvino backend"
-    )
+    element_type = None
+    if isinstance(x1, OpenVINOKerasTensor):
+        element_type = x1.output.get_element_type()
+    if isinstance(x2, OpenVINOKerasTensor):
+        element_type = x2.output.get_element_type()
+    x1 = get_ov_output(x1, element_type)
+    x2 = get_ov_output(x2, element_type)
+    x1, x2 = _align_operand_types(x1, x2, "divide_no_nan()")
+
+    zero = ov_opset.constant(0, x2.get_element_type())
+    div = ov_opset.divide(x1, x2)
+    is_zero = ov_opset.equal(x2, zero)
+    result = ov_opset.select(is_zero, zero, div)
+    return OpenVINOKerasTensor(result.output(0))
 
 
 def true_divide(x1, x2):

{keras_nightly-3.14.0.dev2026012504 → keras_nightly-3.14.0.dev2026012704}/keras/src/backend/torch/export.py

@@ -36,15 +36,31 @@ class TorchExportArchive:
             )
 
         sample_inputs = tree.map_structure(
-            lambda x: convert_spec_to_tensor(x, replace_none_number=1),
+            lambda x: convert_spec_to_tensor(x, replace_none_number=2),
             input_signature,
         )
         sample_inputs = tuple(sample_inputs)
 
+        # Build dynamic_shapes from input_signature where shape has None
+        # Use a shared "batch" dim for dimension 0 across all inputs
+        batch_dim = torch.export.Dim("batch", min=1)
+        dynamic_shapes = []
+        for spec in input_signature:
+            dim_spec = {}
+            for dim_idx, dim_val in enumerate(spec.shape):
+                if dim_val is None:
+                    if dim_idx == 0:
+                        dim_spec[dim_idx] = batch_dim
+                    else:
+                        dim_spec[dim_idx] = torch.export.Dim(
+                            f"dim_{len(dynamic_shapes)}_{dim_idx}", min=1
+                        )
+            dynamic_shapes.append(dim_spec if dim_spec else None)
+        dynamic_shapes = tuple(dynamic_shapes) if any(dynamic_shapes) else None
+
         # Ref: torch_xla.tf_saved_model_integration
-        # TODO: Utilize `dynamic_shapes`
         exported = torch.export.export(
-            resource, sample_inputs, dynamic_shapes=None, strict=False
+            resource, sample_inputs, dynamic_shapes=dynamic_shapes, strict=False
         )
         options = torch_xla.stablehlo.StableHLOExportOptions(
             override_tracing_arguments=sample_inputs

{keras_nightly-3.14.0.dev2026012504 → keras_nightly-3.14.0.dev2026012704}/keras/src/layers/normalization/batch_normalization.py

@@ -8,6 +8,26 @@ from keras.src.layers.input_spec import InputSpec
 from keras.src.layers.layer import Layer
 
 
+# TODO(abheesht17): Move this to utils?
+def _clone_initializer(initializer):
+    """Clones an initializer to ensure a new seed.
+
+    Args:
+        initializer: The initializer to clone.
+
+    Returns:
+        A cloned initializer if it is clonable, otherwise the original one.
+
+    As of tensorflow 2.10, we need to clone user passed initializers when
+    invoking them twice to avoid creating the same randomized initialization.
+    """
+    if isinstance(initializer, initializers.Initializer):
+        config = initializer.get_config()
+        return initializer.__class__.from_config(config)
+    # If we get a string or dict, just return as we cannot and should not clone.
+    return initializer
+
+
 @keras_export("keras.layers.BatchNormalization")
 class BatchNormalization(Layer):
     """Layer that normalizes its inputs.
@@ -72,6 +92,22 @@ class BatchNormalization(Layer):
             variance) for the layer across all devices at each training step
             in a distributed training strategy.
             If `False`, each replica uses its own local batch statistics.
+        renorm: Whether to use
+            [Batch Renormalization](https://arxiv.org/abs/1702.03275). This
+            adds extra variables during training. The inference is the same
+            for either value of this parameter.
+        renorm_clipping: Dictionary, valid only if `renorm = True`.
+            Maps optional keys `"rmax"`, `"rmin"`, `"dmax"` to floats used to
+            clip the renorm correction. The correction `(r, d)` is used as
+            `corrected_value = normalized_value * r + d`, with `r` clipped to
+            `[rmin, rmax]`, and `d` to `[-dmax, dmax]`. Missing `rmax`, `rmin`,
+            `dmax` are set to `inf`, `0`, `inf`, respectively.
+        renorm_momentum: Momentum used to update the moving means and standard
+            deviations with renorm. Valid only if `renorm= True`. Unlike
+            `momentum`, this affects training and should be neither too small
+            (which would add noise) nor too large (which would give stale
+            estimates). Note that `momentum` is still applied to get the means
+            and variances for inference.
         **kwargs: Base layer keyword arguments (e.g. `name` and `dtype`).
 
     Call arguments:
@@ -136,6 +172,9 @@ class BatchNormalization(Layer):
         gamma_regularizer=None,
         beta_constraint=None,
         gamma_constraint=None,
+        renorm=False,
+        renorm_clipping=None,
+        renorm_momentum=0.99,
         synchronized=False,
         **kwargs,
     ):
@@ -165,6 +204,33 @@ class BatchNormalization(Layer):
         self.gamma_constraint = constraints.get(gamma_constraint)
         self.supports_masking = True
 
+        self.renorm = renorm
+        if renorm:
+            renorm_clipping = renorm_clipping or {}
+            keys = ["rmax", "rmin", "dmax"]
+            if set(renorm_clipping) - set(keys):
+                raise ValueError(
+                    "Received invalid keys for `renorm_clipping` argument: "
+                    f"{renorm_clipping}. Supported values: {keys}."
+                )
+            rmax = renorm_clipping.get("rmax")
+            rmin = renorm_clipping.get("rmin")
+            dmax = renorm_clipping.get("dmax")
+
+            if rmax is not None and rmin is not None and rmax < rmin:
+                raise ValueError(
+                    "rmax should be greater than rmin in the `renorm_clipping` "
+                    "argument. Received: rmax={rmax}, rmin={rmin}."
+                )
+            if dmax is not None and dmax < 0:
+                raise ValueError(
+                    "dmax should be non-negative in the `renorm_clipping` "
+                    """argument. Received: dmax={dmax}."""
+                )
+
+        self.renorm_clipping = renorm_clipping
+        self.renorm_momentum = renorm_momentum
+
         self.gamma = None
         self.beta = None
         self.moving_mean = None
@@ -208,6 +274,40 @@ class BatchNormalization(Layer):
             autocast=False,
         )
 
+        if self.renorm:
+            # In batch renormalization we track the inference moving stddev
+            # instead of the moving variance to more closely align with the
+            # paper. The stddev is initialized as sqrt of the variance
+            # initializer.
+            def moving_stddev_initializer(shape, dtype=None):
+                cloned = _clone_initializer(self.moving_variance_initializer)
+                return ops.sqrt(cloned(shape, dtype=dtype))
+
+            self.moving_stddev = self.add_weight(
+                shape=shape,
+                name="moving_stddev",
+                initializer=moving_stddev_initializer,
+                trainable=False,
+                autocast=False,
+            )
+            # Create variables to maintain the moving mean and standard
+            # deviation. These are used in training and thus are different
+            # from the moving averages above.
+            self.renorm_mean = self.add_weight(
+                shape=shape,
+                name="renorm_mean",
+                initializer=_clone_initializer(self.moving_mean_initializer),
+                trainable=False,
+                autocast=False,
+            )
+            self.renorm_stddev = self.add_weight(
+                shape=shape,
+                name="renorm_stddev",
+                initializer=moving_stddev_initializer,
+                trainable=False,
+                autocast=False,
+            )
+
         self.input_spec = InputSpec(
             ndim=len(input_shape), axes={self.axis: input_shape[self.axis]}
         )
@@ -250,20 +350,6 @@ class BatchNormalization(Layer):
         moving_mean = ops.cast(self.moving_mean, inputs.dtype)
         moving_variance = ops.cast(self.moving_variance, inputs.dtype)
 
-        if training and self.trainable:
-            mean, variance = self._moments(inputs, mask)
-
-            self.moving_mean.assign(
-                moving_mean * self.momentum + mean * (1.0 - self.momentum)
-            )
-            self.moving_variance.assign(
-                moving_variance * self.momentum
-                + variance * (1.0 - self.momentum)
-            )
-        else:
-            mean = moving_mean
-            variance = moving_variance
-
         if self.scale:
             gamma = ops.cast(self.gamma, inputs.dtype)
         else:
@@ -274,6 +360,39 @@ class BatchNormalization(Layer):
         else:
             beta = None
 
+        if training and self.trainable:
+            mean, variance = self._moments(inputs, mask)
+
+            if self.renorm:
+                # Compute renorm corrections (r and d).
+                (
+                    r,
+                    d,
+                    mean,
+                    variance,
+                ) = self._renorm_correction_and_moments(mean, variance)
+
+                # x = x * gamma + beta without renorm, and
+                # (x * r + d) * gamma + beta = x * (r * gamma) + (d * gamma +
+                # beta) with renorm.
+                gamma, beta = self._compose_transforms(
+                    r, d, gamma, beta, inputs.dtype
+                )
+
+                # Update moving statistics.
+                self._update_renorm_statistics(mean, variance)
+            else:
+                self.moving_mean.assign(
+                    moving_mean * self.momentum + mean * (1.0 - self.momentum)
+                )
+                self.moving_variance.assign(
+                    moving_variance * self.momentum
+                    + variance * (1.0 - self.momentum)
+                )
+        else:
+            mean = moving_mean
+            variance = moving_variance
+
         outputs = ops.batch_normalization(
             x=inputs,
             mean=mean,
@@ -306,6 +425,9 @@ class BatchNormalization(Layer):
             "beta_constraint": constraints.serialize(self.beta_constraint),
             "gamma_constraint": constraints.serialize(self.gamma_constraint),
             "synchronized": self.synchronized,
+            "renorm": self.renorm,
+            "renorm_clipping": self.renorm_clipping,
+            "renorm_momentum": self.renorm_momentum,
         }
         return {**base_config, **config}
 
@@ -346,3 +468,115 @@ class BatchNormalization(Layer):
         variance = weighted_distsq / (sum_of_weights + backend.epsilon())
 
         return ops.squeeze(mean), ops.squeeze(variance)
+
+    def _renorm_correction_and_moments(self, mean, variance):
+        """Computes the correction for batch renormalization.
+
+        This method computes the r and d correction factors.
+
+        Args:
+            mean: The mean of the current batch.
+            variance: The variance of the current batch.
+
+        Returns:
+            A tuple (r, s, mean, variance) where r and d are the correction
+            factors, and mean/variance are passed through unchanged.
+        """
+        stddev = ops.sqrt(variance + self.epsilon)
+
+        # Get the renorm moving statistics.
+        renorm_mean = ops.cast(self.renorm_mean, mean.dtype)
+        # Avoid divide by zero early on in training.
+        renorm_stddev = ops.maximum(
+            ops.cast(self.renorm_stddev, mean.dtype),
+            ops.sqrt(ops.cast(self.epsilon, mean.dtype)),
+        )
+
+        # Compute the corrections for batch renorm.
+        r = ops.divide(stddev, renorm_stddev)
+        d = ops.divide(ops.subtract(mean, renorm_mean), renorm_stddev)
+
+        # Apply clipping.
+        rmin = self.renorm_clipping.get("rmin")
+        rmax = self.renorm_clipping.get("rmax")
+        dmax = self.renorm_clipping.get("dmax")
+
+        if rmin is not None:
+            r = ops.maximum(r, rmin)
+        if rmax is not None:
+            r = ops.minimum(r, rmax)
+        if dmax is not None:
+            d = ops.clip(d, -dmax, dmax)
+
+        return r, d, mean, variance
+
+    def _compose_transforms(self, r, d, gamma, beta, dtype):
+        """Composes the renorm correction with gamma and beta.
+
+        When training with renorm, the normalized values (x) are transformed
+        as: (x * r + d) * gamma + beta = x * (r * gamma) + (d * gamma + beta).
+        This method computes the effective scale and offset.
+
+        Args:
+            r: The r correction factor.
+            d: The d correction factor.
+            gamma: The gamma (scale) parameter, or None.
+            beta: The beta (offset) parameter, or None.
+            dtype: The dtype for the output.
+
+        Returns:
+            A tuple (effective_gamma, effective_beta).
+        """
+        r = ops.stop_gradient(r)
+        d = ops.stop_gradient(d)
+
+        if gamma is not None:
+            effective_gamma = ops.multiply(r, gamma)
+            effective_beta = ops.multiply(d, gamma)
+        else:
+            effective_gamma = ops.cast(r, dtype)
+            effective_beta = ops.cast(d, dtype)
+
+        if beta is not None:
+            effective_beta = ops.add(effective_beta, beta)
+
+        return effective_gamma, effective_beta
+
+    def _update_renorm_statistics(self, mean, variance):
+        """Updates the renorm and moving statistics.
+        Args:
+            mean: The mean of the current batch.
+            variance: The variance of the current batch.
+        """
+        stddev = ops.sqrt(variance + self.epsilon)
+
+        # Update renorm moving mean and stddev.
+        renorm_mean = ops.cast(self.renorm_mean, mean.dtype)
+        renorm_stddev = ops.cast(self.renorm_stddev, mean.dtype)
+
+        self.renorm_mean.assign(
+            renorm_mean * self.renorm_momentum
+            + mean * (1.0 - self.renorm_momentum)
+        )
+        self.renorm_stddev.assign(
+            renorm_stddev * self.renorm_momentum
+            + stddev * (1.0 - self.renorm_momentum)
+        )
+
+        moving_mean = ops.cast(self.moving_mean, mean.dtype)
+        moving_stddev = ops.cast(self.moving_stddev, mean.dtype)
+
+        self.moving_mean.assign(
+            moving_mean * self.momentum + mean * (1.0 - self.momentum)
+        )
+
+        new_moving_stddev = moving_stddev * self.momentum + stddev * (
+            1.0 - self.momentum
+        )
+        self.moving_stddev.assign(new_moving_stddev)
+
+        # Derive `moving_variance` from `moving_stddev`, applying ReLU in case
+        # floating point rounding causes it to go negative.
+        self.moving_variance.assign(
+            ops.relu(new_moving_stddev * new_moving_stddev - self.epsilon)
+        )

{keras_nightly-3.14.0.dev2026012504 → keras_nightly-3.14.0.dev2026012704}/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.dev2026012504"
+__version__ = "3.14.0.dev2026012704"
 
 
 @keras_export("keras.version")

{keras_nightly-3.14.0.dev2026012504 → keras_nightly-3.14.0.dev2026012704}/keras_nightly.egg-info/PKG-INFO

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