Merge branch 'master' into gsoc2025

Mohamed-Ashraf273 · Mohamed-Ashraf273 · commit c278dbf1cfa0 · 2025-07-18T22:09:23.000+03:00
diff --git a/keras/api/_tf_keras/keras/ops/__init__.py b/keras/api/_tf_keras/keras/ops/__init__.py
@@ -191,6 +191,7 @@
 from keras.src.ops.numpy import greater_equal as greater_equal
 from keras.src.ops.numpy import hamming as hamming
 from keras.src.ops.numpy import hanning as hanning
+from keras.src.ops.numpy import heaviside as heaviside
 from keras.src.ops.numpy import histogram as histogram
 from keras.src.ops.numpy import hstack as hstack
 from keras.src.ops.numpy import identity as identity
diff --git a/keras/api/_tf_keras/keras/ops/numpy/__init__.py b/keras/api/_tf_keras/keras/ops/numpy/__init__.py
@@ -80,6 +80,7 @@
 from keras.src.ops.numpy import greater_equal as greater_equal
 from keras.src.ops.numpy import hamming as hamming
 from keras.src.ops.numpy import hanning as hanning
+from keras.src.ops.numpy import heaviside as heaviside
 from keras.src.ops.numpy import histogram as histogram
 from keras.src.ops.numpy import hstack as hstack
 from keras.src.ops.numpy import identity as identity
diff --git a/keras/api/ops/__init__.py b/keras/api/ops/__init__.py
@@ -191,6 +191,7 @@
 from keras.src.ops.numpy import greater_equal as greater_equal
 from keras.src.ops.numpy import hamming as hamming
 from keras.src.ops.numpy import hanning as hanning
+from keras.src.ops.numpy import heaviside as heaviside
 from keras.src.ops.numpy import histogram as histogram
 from keras.src.ops.numpy import hstack as hstack
 from keras.src.ops.numpy import identity as identity
diff --git a/keras/api/ops/numpy/__init__.py b/keras/api/ops/numpy/__init__.py
@@ -80,6 +80,7 @@
 from keras.src.ops.numpy import greater_equal as greater_equal
 from keras.src.ops.numpy import hamming as hamming
 from keras.src.ops.numpy import hanning as hanning
+from keras.src.ops.numpy import heaviside as heaviside
 from keras.src.ops.numpy import histogram as histogram
 from keras.src.ops.numpy import hstack as hstack
 from keras.src.ops.numpy import identity as identity
diff --git a/keras/src/backend/jax/numpy.py b/keras/src/backend/jax/numpy.py
@@ -52,6 +52,12 @@ def hanning(x):
     return jnp.hanning(x)
 
 
+def heaviside(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    return jnp.heaviside(x1, x2)
+
+
 def kaiser(x, beta):
     x = convert_to_tensor(x)
     return jnp.kaiser(x, beta)
diff --git a/keras/src/backend/numpy/numpy.py b/keras/src/backend/numpy/numpy.py
@@ -320,6 +320,19 @@ def hanning(x):
     return np.hanning(x).astype(config.floatx())
 
 
+def heaviside(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    if dtype in ["int8", "int16", "int32", "uint8", "uint16", "uint32"]:
+        dtype = config.floatx()
+    elif dtype in ["int64"]:
+        dtype = "float64"
+
+    return np.heaviside(x1, x2).astype(dtype)
+
+
 def kaiser(x, beta):
     x = convert_to_tensor(x)
     return np.kaiser(x, beta).astype(config.floatx())
diff --git a/keras/src/backend/openvino/excluded_concrete_tests.txt b/keras/src/backend/openvino/excluded_concrete_tests.txt
@@ -11,6 +11,7 @@ NumpyDtypeTest::test_bartlett
 NumpyDtypeTest::test_blackman
 NumpyDtypeTest::test_hamming
 NumpyDtypeTest::test_hanning
+NumpyDtypeTest::test_heaviside
 NumpyDtypeTest::test_kaiser
 NumpyDtypeTest::test_bitwise
 NumpyDtypeTest::test_cbrt
@@ -137,6 +138,7 @@ NumpyTwoInputOpsCorrectnessTest::test_cross
 NumpyTwoInputOpsCorrectnessTest::test_digitize
 NumpyTwoInputOpsCorrectnessTest::test_divide_no_nan
 NumpyTwoInputOpsCorrectnessTest::test_einsum
+NumpyTwoInputOpsCorrectnessTest::test_heaviside
 NumpyTwoInputOpsCorrectnessTest::test_inner
 NumpyTwoInputOpsCorrectnessTest::test_linspace
 NumpyTwoInputOpsCorrectnessTest::test_logspace
@@ -155,6 +157,8 @@ NumpyOneInputOpsDynamicShapeTest::test_kaiser
 NumpyOneInputOpsStaticShapeTest::test_angle
 NumpyOneInputOpsStaticShapeTest::test_cbrt
 NumpyOneInputOpsStaticShapeTest::test_deg2rad
+NumpyTwoInputOpsDynamicShapeTest::test_heaviside
+NumpyTwoInputOpsStaticShapeTest::test_heaviside
 CoreOpsBehaviorTests::test_associative_scan_invalid_arguments
 CoreOpsBehaviorTests::test_scan_invalid_arguments
 CoreOpsCallsTests::test_associative_scan_basic_call
@@ -195,7 +199,6 @@ MathOpsCorrectnessTest::test_istft4
 MathOpsCorrectnessTest::test_istft5
 MathOpsCorrectnessTest::test_istft6
 MathOpsCorrectnessTest::test_logdet
-MathOpsCorrectnessTest::test_logsumexp
 MathOpsCorrectnessTest::test_rfft0
 MathOpsCorrectnessTest::test_rfft1
 MathOpsCorrectnessTest::test_rfft2
@@ -277,4 +280,4 @@ TestMathErrors::test_invalid_fft_length
 TestMathErrors::test_istft_invalid_window_shape_2D_inputs
 TestMathErrors::test_stft_invalid_input_type
 TestMathErrors::test_stft_invalid_window
-TestMathErrors::test_stft_invalid_window_shape
+TestMathErrors::test_stft_invalid_window_shape
diff --git a/keras/src/backend/openvino/math.py b/keras/src/backend/openvino/math.py
@@ -44,13 +44,17 @@ def logsumexp(x, axis=None, keepdims=False):
         axis = list(axis)
     axis = ov_opset.constant(axis, Type.i32).output(0)
     const_zero = ov_opset.constant(0, x.get_element_type()).output(0)
-    reduce_max = ov_opset.reduce_max(x, axis, keepdims).output(0)
+    # Use keepdims=True for reduce_max to ensure proper broadcasting
+    reduce_max = ov_opset.reduce_max(x, axis, True).output(0)
     is_finite = ov_opset.is_finite(reduce_max).output(0)
     norm_max = ov_opset.select(is_finite, reduce_max, const_zero).output(0)
     norm_max_sub = ov_opset.subtract(x, norm_max).output(0)
     exp_norm_max = ov_opset.exp(norm_max_sub).output(0)
     sum_exp = ov_opset.reduce_sum(exp_norm_max, axis, keepdims).output(0)
     log_sum_exp = ov_opset.log(sum_exp).output(0)
+    # Squeeze norm_max if needed to match dimensions
+    if not keepdims:
+        norm_max = ov_opset.squeeze(norm_max, axis).output(0)
     log_sum_exp = ov_opset.add(norm_max, log_sum_exp).output(0)
     return OpenVINOKerasTensor(log_sum_exp)
 
diff --git a/keras/src/backend/openvino/numpy.py b/keras/src/backend/openvino/numpy.py
@@ -481,6 +481,12 @@ def hamming(x):
     )
 
 
+def heaviside(x1, x2):
+    raise NotImplementedError(
+        "`heaviside` is not supported with openvino backend"
+    )
+
+
 def kaiser(x, beta):
     raise NotImplementedError("`kaiser` is not supported with openvino backend")
 
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
@@ -156,6 +156,26 @@ def hanning(x):
     return tf.signal.hann_window(x, periodic=False)
 
 
+def heaviside(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    if dtype in ["int8", "int16", "int32", "uint8", "uint16", "uint32"]:
+        dtype = config.floatx()
+    elif dtype in ["int64"]:
+        dtype = "float64"
+
+    x1 = tf.cast(x1, dtype)
+    x2 = tf.cast(x2, dtype)
+
+    return tf.where(
+        x1 < 0,
+        tf.zeros_like(x1),
+        tf.where(x1 > 0, tf.ones_like(x1), x2),
+    )
+
+
 def kaiser(x, beta):
     x = convert_to_tensor(x, dtype=tf.int32)
     return tf.signal.kaiser_window(x, beta=beta)
diff --git a/keras/src/backend/torch/numpy.py b/keras/src/backend/torch/numpy.py
@@ -445,6 +445,22 @@ def hanning(x):
     return torch.signal.windows.hann(x)
 
 
+def heaviside(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    if dtype in ["int8", "int16", "int32", "uint8", "uint16", "uint32"]:
+        dtype = config.floatx()
+    elif dtype == "int64":
+        dtype = "float64"
+
+    x1 = cast(x1, dtype)
+    x2 = cast(x2, dtype)
+
+    return torch.heaviside(x1, x2)
+
+
 def kaiser(x, beta):
     x = convert_to_tensor(x)
     return torch.signal.windows.kaiser(x, beta=beta)
diff --git a/keras/src/layers/core/dense.py b/keras/src/layers/core/dense.py
@@ -693,53 +693,56 @@ def _get_kernel_with_merged_lora(self):
                 `kernel_scale`: The quantization scale for the merged kernel.
                     This is `None` if the layer is not quantized.
         """
-        if self.dtype_policy.quantization_mode is not None:
-            kernel_value = self._kernel
-            kernel_scale = self.kernel_scale
-            if self.lora_enabled:
-                # Dequantize kernel to float
-                if self.quantization_mode == "int4":
-                    unpacked_kernel = quantizers.unpack_int4(
-                        kernel_value, self._orig_input_dim
-                    )
-                    float_kernel = ops.divide(
-                        ops.cast(unpacked_kernel, self.compute_dtype),
-                        kernel_scale,
-                    )
-                    quant_range = (-8, 7)
-                elif self.quantization_mode == "int8":
-                    float_kernel = ops.divide(
-                        ops.cast(kernel_value, self.compute_dtype), kernel_scale
-                    )
-                    quant_range = (-127, 127)
-                else:
-                    raise ValueError(
-                        "Unsupported quantization mode: "
-                        f"{self.quantization_mode}"
-                    )
-
-                # Merge LoRA weights in float domain
-                lora_delta = (self.lora_alpha / self.lora_rank) * ops.matmul(
-                    self.lora_kernel_a, self.lora_kernel_b
-                )
-                merged_float_kernel = ops.add(float_kernel, lora_delta)
-
-                # Requantize
-                requantized_kernel, kernel_scale = quantizers.abs_max_quantize(
-                    merged_float_kernel,
-                    axis=0,
-                    value_range=quant_range,
-                    dtype="int8",
-                    to_numpy=True,
-                )
-                kernel_scale = ops.squeeze(kernel_scale, axis=0)
-
-                # Pack if int4
-                if self.quantization_mode == "int4":
-                    kernel_value, _, _ = quantizers.pack_int4(
-                        requantized_kernel
-                    )
-                else:
-                    kernel_value = requantized_kernel
+        if self.dtype_policy.quantization_mode is None:
+            return self.kernel, None
+
+        kernel_value = self._kernel
+        kernel_scale = self.kernel_scale
+
+        if not self.lora_enabled:
             return kernel_value, kernel_scale
-        return self.kernel, None
+
+        # Dequantize, Merge, and Re-quantize
+
+        # Dequantize kernel to float
+        if self.quantization_mode == "int4":
+            unpacked_kernel = quantizers.unpack_int4(
+                kernel_value, self._orig_input_dim
+            )
+            float_kernel = ops.divide(
+                ops.cast(unpacked_kernel, self.compute_dtype),
+                kernel_scale,
+            )
+            quant_range = (-8, 7)
+        elif self.quantization_mode == "int8":
+            float_kernel = ops.divide(
+                ops.cast(kernel_value, self.compute_dtype), kernel_scale
+            )
+            quant_range = (-127, 127)
+        else:
+            raise ValueError(
+                f"Unsupported quantization mode: {self.quantization_mode}"
+            )
+
+        # Merge LoRA weights in float domain
+        lora_delta = (self.lora_alpha / self.lora_rank) * ops.matmul(
+            self.lora_kernel_a, self.lora_kernel_b
+        )
+        merged_float_kernel = ops.add(float_kernel, lora_delta)
+
+        # Requantize
+        requantized_kernel, kernel_scale = quantizers.abs_max_quantize(
+            merged_float_kernel,
+            axis=0,
+            value_range=quant_range,
+            dtype="int8",
+            to_numpy=True,
+        )
+        kernel_scale = ops.squeeze(kernel_scale, axis=0)
+
+        # Pack if int4
+        if self.quantization_mode == "int4":
+            kernel_value, _, _ = quantizers.pack_int4(requantized_kernel)
+        else:
+            kernel_value = requantized_kernel
+        return kernel_value, kernel_scale
diff --git a/keras/src/layers/core/einsum_dense.py b/keras/src/layers/core/einsum_dense.py
diff --git a/keras/src/layers/core/einsum_dense_test.py b/keras/src/layers/core/einsum_dense_test.py
diff --git a/keras/src/losses/losses.py b/keras/src/losses/losses.py
diff --git a/keras/src/ops/numpy.py b/keras/src/ops/numpy.py
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
