diff --git a/dm_pix/_src/patch.py b/dm_pix/_src/patch.py
index 2d9f593..ea2f563 100644
--- a/dm_pix/_src/patch.py
+++ b/dm_pix/_src/patch.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """This module provides image patching functionality."""
 
+import math
 from typing import Sequence
 
 import chex
@@ -89,7 +90,8 @@ def extract_patches(
 
   # `conv_general_dilated_patches` returns `patches` in channel-major order,
   # rearrange to match interface of `tf.image.extract_patches`.
-  patches = jnp.reshape(patches, patches.shape[:-1] + (channels, -1))
+  output_shape = patches.shape
+  patch_size = math.prod(sizes[1:-1])
+  patches = jnp.reshape(patches, output_shape[:-1] + (channels, patch_size))
   patches = jnp.moveaxis(patches, -2, -1)
-  patches = jnp.reshape(patches, patches.shape[:-2] + (-1,))
-  return patches
+  return jnp.reshape(patches, output_shape)
diff --git a/dm_pix/_src/patch_test.py b/dm_pix/_src/patch_test.py
index 0ad9bd7..b45ab57 100644
--- a/dm_pix/_src/patch_test.py
+++ b/dm_pix/_src/patch_test.py
@@ -30,12 +30,11 @@ def _create_test_images(shape):
 class PatchTest(chex.TestCase, parameterized.TestCase):
 
   @chex.all_variants
-  @parameterized.named_parameters(
-      ('padding_valid', 'VALID'),
-      ('padding_same', 'SAME'),
+  @parameterized.product(
+      padding=('VALID', 'SAME'),
+      image_shape=((2, 5, 7, 3), (0, 5, 7, 3), (2, 2, 7, 3)),
   )
-  def test_extract_patches(self, padding):
-    image_shape = (2, 5, 7, 3)
+  def test_extract_patches(self, padding, image_shape):
     images = _create_test_images(image_shape)
 
     sizes = (1, 2, 3, 1)