Repository: keras-team/keras

Issue:
Missing block_size validation in space_to_depth and depth_to_space leading to DivisionByZero


## Summary

The recently added `space_to_depth` and `depth_to_space` operations in `keras.ops.nn` do **not** validate the `block_size` argument.

Providing `block_size=0` or a negative value results in:

- Backend-specific errors (TensorFlow raises a low-level op error)
- In some cases, Python-level `ZeroDivisionError` during shape inference
- Inconsistent behavior across backends (TF / JAX / Torch)

Keras should raise a consistent, frontend-level `ValueError` before dispatching to backends, similar to how other ops (e.g., GLU axis validation) are handled.

---

## Reproduction Script

```python
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import numpy as np

# Input shape: (batch, height, width, channels)
x = np.random.rand(1, 4, 4, 12).astype("float32")

# This should raise a Keras-level ValueError,
# but currently the backend raises the exception.
try:
    out = keras.ops.nn.space_to_depth(x, block_size=0)
except Exception as e:
    print(f"Caught: {e}")
```

### Actual Behavior (TensorFlow backend)
The error is raised by the TensorFlow runtime after the op is dispatched:
```
InvalidArgumentError: Value for attr 'block_size' of 0 must be at least minimum 2
```
### Expected Behavior (Keras Frontend)
Keras should validate this at the API level (in `nn.py`) and raise a consistent `ValueError` before reaching the backend:

```
ValueError: block_size must be >= 2. Received: block_size=0
```
---

## Root Cause

In `keras/src/ops/nn.py`, the `SpaceToDepth` and `DepthToSpace` classes:

- Pass `block_size` directly to backend ops without validating it
- Do not mirror TensorFlow’s minimum constraint (`block_size >= 2`)
- Do not guard against zero or negative values at the API level

This results in backend-dependent behavior and violates Keras' expected consistency across frameworks.

---

## Suggested Fix

Add a simple validation before backend dispatch inside the op class (constructor or call method):

```python
if block_size < 2:
    raise ValueError(
        f"`block_size` must be at least 2. Received: block_size={block_size}"
    )
```

---

## Why This Matters

- Ensures **backend-agnostic correctness** across TF / JAX / Torch.
- Prevents confusing low-level backend exceptions.
- Aligns with Keras' design principles (“fail early, fail consistently”).
- Maintains symmetry with other ops that enforce input constraints.

---

## Environment

- OS: Windows 11
- Python: 3.13.5
- Keras: 3.14.0
- Backend: TensorFlow

---

## Conclusion

This is a small but important correctness issue.  
Adding frontend validation ensures consistent and predictable behavior for all Keras backends and users.

Discussion hints:

Solution patch:
diff --git a/keras/src/ops/nn.py b/keras/src/ops/nn.py
index 29baa5d56afe..27d23c8a2923 100644
--- a/keras/src/ops/nn.py
+++ b/keras/src/ops/nn.py
@@ -3402,6 +3402,11 @@ def depth_to_space(x, block_size, data_format="channels_last"):
     (1, 3, 4, 4)
     """
     data_format = standardize_data_format(data_format)
+    if block_size < 2:
+        raise ValueError(
+            "`block_size` must be at least 2. "
+            f"Received: block_size={block_size}"
+        )
     if any_symbolic_tensors((x,)):
         return DepthToSpace(block_size, data_format=data_format).symbolic_call(
             x
@@ -3501,6 +3506,11 @@ def space_to_depth(x, block_size, data_format="channels_last"):
     (1, 12, 2, 2)
     """
     data_format = standardize_data_format(data_format)
+    if block_size < 2:
+        raise ValueError(
+            "`block_size` must be at least 2. "
+            f"Received: block_size={block_size}"
+        )
     if any_symbolic_tensors((x,)):
         return SpaceToDepth(block_size, data_format=data_format).symbolic_call(
             x


Test patch:
diff --git a/keras/src/ops/nn_test.py b/keras/src/ops/nn_test.py
index 86c01492b430..bab099f3577e 100644
--- a/keras/src/ops/nn_test.py
+++ b/keras/src/ops/nn_test.py
@@ -3600,3 +3600,43 @@ def test_depth_to_space_space_to_depth_roundtrip(self):
         y = knn.depth_to_space(x, block_size=2, data_format="channels_first")
         z = knn.space_to_depth(y, block_size=2, data_format="channels_first")
         self.assertAllClose(x, z)
+
+    def test_depth_to_space_block_size_validation(self):
+        x = ops.arange(48, dtype="float32")
+        x = ops.reshape(x, [1, 2, 2, 12])
+
+        # block_size must be at least 2
+        with self.assertRaisesRegex(
+            ValueError, "`block_size` must be at least 2"
+        ):
+            knn.depth_to_space(x, block_size=0)
+
+        with self.assertRaisesRegex(
+            ValueError, "`block_size` must be at least 2"
+        ):
+            knn.depth_to_space(x, block_size=1)
+
+        with self.assertRaisesRegex(
+            ValueError, "`block_size` must be at least 2"
+        ):
+            knn.depth_to_space(x, block_size=-1)
+
+    def test_space_to_depth_block_size_validation(self):
+        x = ops.arange(48, dtype="float32")
+        x = ops.reshape(x, [1, 4, 4, 3])
+
+        # block_size must be at least 2
+        with self.assertRaisesRegex(
+            ValueError, "`block_size` must be at least 2"
+        ):
+            knn.space_to_depth(x, block_size=0)
+
+        with self.assertRaisesRegex(
+            ValueError, "`block_size` must be at least 2"
+        ):
+            knn.space_to_depth(x, block_size=1)
+
+        with self.assertRaisesRegex(
+            ValueError, "`block_size` must be at least 2"
+        ):
+            knn.space_to_depth(x, block_size=-1)