Fix a bug where SSIM fails on non-square images

This PR also adds unit tests that probe non-square images.

PiperOrigin-RevId: 401296829

dm_pix/_src/metrics.py

@@ -187,7 +187,7 @@ def filt_fn_y(z):
   def filt_fn_x(z):
     z_flat = jnp.moveaxis(z, -2, -1).reshape((-1, z.shape[-2]))
     z_filt_shape = ((z.shape[-4],) if z.ndim == 4 else
-                    ()) + (z.shape[-2], z.shape[-1], -1)
+                    ()) + (z.shape[-3], z.shape[-1], -1)
     return jnp.moveaxis(filt_fn_vmap(z_flat).reshape(z_filt_shape), -1, -2)
 
   # Apply the blur in both x and y.

dm_pix/_src/metrics_test.py

@@ -64,7 +64,7 @@ def test_ssim_golden(self):
     """Test that the SSIM implementation matches the Tensorflow version."""
 
     key = jax.random.PRNGKey(0)
-    for shape in (2, 12, 12, 3), (12, 12, 3):
+    for shape in ((2, 12, 12, 3), (12, 12, 3), (2, 12, 15, 3), (17, 12, 3)):
       for _ in range(4):
         (max_val_key, img0_key, img1_key, filter_size_key, filter_sigma_key,
          k1_key, k2_key, key) = jax.random.split(key, 8)
@@ -99,9 +99,13 @@ def test_ssim_golden(self):
               ))
           ssim = ssim_fn(img0, img1)
           if not return_map:
-            self.assertAllClose(ssim, ssim_gt)
+            self.assertAllClose(ssim, ssim_gt, atol=1e-5, rtol=1e-5)
           else:
-            self.assertAllClose(np.mean(ssim, list(range(-3, 0))), ssim_gt)
+            self.assertAllClose(
+                np.mean(ssim, list(range(-3, 0))),
+                ssim_gt,
+                atol=1e-5,
+                rtol=1e-5)
           self.assertLessEqual(np.max(ssim), 1.)
           self.assertGreaterEqual(np.min(ssim), -1.)