Address bug with convolution using Tensorflow, Numpy, Jax backends (#21796)

* add exception

* Update keras/src/backend/tensorflow/nn.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update keras/src/layers/convolutional/conv_test.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* address numpy + make test more generic

* fix jax

* fix pydocs

* fix error msg in tensorflow

* handle only static case

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: danielenricocahall

keras/src/backend/jax/nn.py

@@ -355,7 +355,7 @@ def conv(
     feature_group_count = channels // kernel_in_channels
     kernel = convert_to_tensor(kernel)
     inputs = convert_to_tensor(inputs, dtype=kernel.dtype)
-    return jax.lax.conv_general_dilated(
+    result = jax.lax.conv_general_dilated(
         inputs,
         kernel,
         strides,
@@ -364,6 +364,14 @@ def conv(
         dimension_numbers=dimension_numbers,
         feature_group_count=feature_group_count,
     )
+    if result.size == 0:
+        raise ValueError(
+            "The convolution operation resulted in an empty output. "
+            "This can happen if the input is too small for the given "
+            "kernel size, strides, dilation rate, and padding mode. "
+            "Please check the input shape and convolution parameters."
+        )
+    return result
 
 
 def depthwise_conv(

keras/src/backend/numpy/nn.py

@@ -404,7 +404,7 @@ def conv(
             f"kernel in_channels {kernel_in_channels}. "
         )
     feature_group_count = channels // kernel_in_channels
-    return np.array(
+    result = np.array(
         jax.lax.conv_general_dilated(
             inputs,
             kernel if is_tensor(kernel) else kernel.numpy(),
@@ -415,6 +415,14 @@ def conv(
             feature_group_count=feature_group_count,
         )
     )
+    if result.size == 0:
+        raise ValueError(
+            "The convolution operation resulted in an empty output. "
+            "This can happen if the input is too small for the given "
+            "kernel size, strides, dilation rate, and padding mode. "
+            "Please check the input shape and convolution parameters."
+        )
+    return result
 
 
 def depthwise_conv(

keras/src/backend/tensorflow/nn.py

@@ -310,14 +310,28 @@ def conv(
 ):
     def _conv():
         tf_data_format = _convert_data_format(data_format, len(inputs.shape))
-        return tf.nn.convolution(
+        result = tf.nn.convolution(
             inputs,
             kernel,
             strides,
             padding.upper(),
             data_format=tf_data_format,
             dilations=dilation_rate,
         )
+        result_shape = result.shape
+        if (
+            result_shape.is_fully_defined()
+            and math.prod(result_shape.as_list()) == 0
+        ):
+            raise ValueError(
+                "The convolution operation resulted in an empty output. "
+                "Output shape:"
+                f" {result_shape}. This can happen if the input is too small "
+                "for the given kernel size, strides, dilation rate, and "
+                "padding mode. Please check the input shape and convolution "
+                "parameters."
+            )
+        return result
 
     # Certain ops are are broken in Tensorflow on CPU only.
     # We can work around by compiling the op with XLA.

keras/src/layers/convolutional/conv_test.py

@@ -1095,3 +1095,11 @@ def test_conv_constraints(self):
         )
         layer.build((None, 5, 5, 3))
         self.assertIsInstance(layer.bias.constraint, constraints.NonNeg)
+
+    def test_conv_raises_exception_on_zero_dims(self):
+        x = np.random.rand(3, 4, 4, 4)
+        l = layers.Conv2D(6, [5, 5], 1, "valid")
+        # The exception type can vary across backends (e.g., ValueError,
+        # tf.errors.InvalidArgumentError, RuntimeError).
+        with self.assertRaises(Exception):
+            l(x)