Fix quantized CDF shape problem.

Author: Johannes Ballé

tensorflow_compression/python/layers/entropy_models.py

@@ -359,8 +359,17 @@ def quantiles_initializer(shape, dtype=None, partition_info=None):
 
     cdf = coder_ops.pmf_to_quantized_cdf(
         pmf, precision=self.range_coder_precision)
+
+    # We need to supply an initializer without fully defined static shape here,
+    # or the variable will return the wrong dynamic shape later. A placeholder
+    # with default gets the trick done.
+    def cdf_init(*args, **kwargs):
+      return array_ops.placeholder_with_default(
+          array_ops.zeros((channels, 1), dtype=dtypes.int32),
+          shape=(channels, None))
+
     self._quantized_cdf = self.add_variable(
-        "quantized_cdf", shape=(channels, 1), dtype=dtypes.int32,
+        "quantized_cdf", shape=None, initializer=cdf_init, dtype=dtypes.int32,
         trainable=False)
 
     update_op = state_ops.assign(

tensorflow_compression/python/layers/entropy_models_test.py

@@ -288,7 +288,7 @@ def test_normalization(self):
       likelihood, = sess.run([likelihood], {inputs: x})
       self.assertEqual(x.shape, likelihood.shape)
       integral = np.sum(likelihood) * .001
-      self.assertAllClose(1, integral, rtol=0, atol=1e-4)
+      self.assertAllClose(1, integral, rtol=0, atol=2e-4)
 
   def test_entropy_estimates(self):
     # Test that entropy estimates match actual range coding.