Fixes example models to match the quantization strategy used in the papers.

PiperOrigin-RevId: 424439812
Change-Id: Iadc25866c4d6a5a769d4c9540b0d2f3123c94f4c

Author: Johannes Ballé

models/bmshj2018.py

@@ -235,7 +235,7 @@ def compress(self, x):
     x_shape = tf.shape(x)[1:-1]
     y_shape = tf.shape(y)[1:-1]
     z_shape = tf.shape(z)[1:-1]
-    z_hat, _ = self.side_entropy_model(z, training=False)
+    z_hat = self.side_entropy_model.quantize(z)
     indexes = self.hyper_synthesis_transform(z_hat)
     indexes = indexes[:, :y_shape[0], :y_shape[1], :]
     side_string = self.side_entropy_model.compress(z)

models/ms2020.py

@@ -209,7 +209,8 @@ def call(self, x, training):
 
     # Build the entropy model for the hyperprior (z).
     em_z = tfc.ContinuousBatchedEntropyModel(
-        self.hyperprior, coding_rank=3, compression=False)
+        self.hyperprior, coding_rank=3, compression=False,
+        offset_heuristic=False)
 
     # When training, z_bpp is based on the noisy version of z (z_tilde).
     _, z_bits = em_z(z, training=training)
@@ -255,7 +256,7 @@ def call(self, x, training):
 
       # For the synthesis transform, use rounding. Note that quantize()
       # overrides the gradient to create a straight-through estimator.
-      y_hat_slice = em_y.quantize(y_slice, sigma, loc=mu)
+      y_hat_slice = em_y.quantize(y_slice, loc=mu)
 
       # Add latent residual prediction (LRP).
       lrp_support = tf.concat([mean_support, y_hat_slice], axis=-1)
@@ -318,7 +319,8 @@ def fit(self, *args, **kwargs):
     retval = super().fit(*args, **kwargs)
     # After training, fix range coding tables.
     self.em_z = tfc.ContinuousBatchedEntropyModel(
-        self.hyperprior, coding_rank=3, compression=True)
+        self.hyperprior, coding_rank=3, compression=True,
+        offset_heuristic=False)
     self.em_y = tfc.LocationScaleIndexedEntropyModel(
         tfc.NoisyNormal, num_scales=self.num_scales, scale_fn=self.scale_fn,
         coding_rank=3, compression=True)