def test_ssim_golden(self):
    """Test our SSIM implementation against the Tensorflow version."""
    rng = random.PRNGKey(0)
    shape = (2, 12, 12, 3)
    for _ in range(4):
      rng, key = random.split(rng)
      max_val = random.uniform(key, minval=0.1, maxval=3.)
      rng, key = random.split(rng)
      img0 = max_val * random.uniform(key, shape=shape, minval=-1, maxval=1)
      rng, key = random.split(rng)
      img1 = max_val * random.uniform(key, shape=shape, minval=-1, maxval=1)
      rng, key = random.split(rng)
      filter_size = random.randint(key, shape=(), minval=1, maxval=10)
      rng, key = random.split(rng)
      filter_sigma = random.uniform(key, shape=(), minval=0.1, maxval=10.)
      rng, key = random.split(rng)
      k1 = random.uniform(key, shape=(), minval=0.001, maxval=0.1)
      rng, key = random.split(rng)
      k2 = random.uniform(key, shape=(), minval=0.001, maxval=0.1)

      ssim_gt = tf.image.ssim(
          img0,
          img1,
          max_val,
          filter_size=filter_size,
          filter_sigma=filter_sigma,
          k1=k1,
          k2=k2).numpy()
      for return_map in [False, True]:
        ssim_fn = jax.jit(
            functools.partial(
                math.compute_ssim,
                max_val=max_val,
                filter_size=filter_size,
                filter_sigma=filter_sigma,
                k1=k1,
                k2=k2,
                return_map=return_map))
        ssim = ssim_fn(img0, img1)
        if not return_map:
          self.assertAllClose(ssim, ssim_gt)
        else:
          self.assertAllClose(np.mean(ssim, [1, 2, 3]), ssim_gt)
        self.assertLessEqual(np.max(ssim), 1.)
        self.assertGreaterEqual(np.min(ssim), -1.)
Test our SSIM implementation against the Tensorflow version.
internal/math_test.py