Fixes tfp.Distribution._parameter_properties inheritance warning.

PiperOrigin-RevId: 424338195
Change-Id: If363b686b74593e93082e0a858e43ffd1263315a

Author: Johannes Ballé

tensorflow_compression/python/distributions/deep_factorized.py

@@ -254,6 +254,11 @@ def _upper_tail(self, tail_mass):
     return helpers.estimate_tails(
         self._logits_cumulative, logits, self.batch_shape_tensor(), self.dtype)
 
+  @classmethod
+  def _parameter_properties(cls, dtype=tf.float32, num_classes=None):
+    raise NotImplementedError(
+        f"`{cls.__name__}` does not implement `_parameter_properties`.")
+
 
 class NoisyDeepFactorized(uniform_noise.UniformNoiseAdapter):
   """DeepFactorized that is convolved with uniform noise."""

tensorflow_compression/python/distributions/round_adapters.py

@@ -157,6 +157,11 @@ def _upper_tail(self, tail_mass):
     return self.transform(helpers.upper_tail(self.base, tail_mass))
   # pylint: enable=protected-access
 
+  @classmethod
+  def _parameter_properties(cls, dtype=tf.float32, num_classes=None):
+    raise NotImplementedError(
+        f"`{cls.__name__}` does not implement `_parameter_properties`.")
+
 
 class RoundAdapter(MonotonicAdapter):
   """Continuous density function + round."""

tensorflow_compression/python/distributions/uniform_noise.py

@@ -189,6 +189,11 @@ def _lower_tail(self, tail_mass):
   def _upper_tail(self, tail_mass):
     return helpers.upper_tail(self.base, tail_mass)
 
+  @classmethod
+  def _parameter_properties(cls, dtype=tf.float32, num_classes=None):
+    raise NotImplementedError(
+        f"`{cls.__name__}` does not implement `_parameter_properties`.")
+
 
 class NoisyMixtureSameFamily(tfp.distributions.MixtureSameFamily):
   """Mixture of distributions with additive i.i.d. uniform noise."""
@@ -211,6 +216,18 @@ def base(self):
     """The base distribution (without uniform noise)."""
     return self._base
 
+  def _batch_shape_tensor(self):
+    return self.base.batch_shape_tensor()
+
+  def _batch_shape(self):
+    return self.base.batch_shape
+
+  def _event_shape_tensor(self):
+    return self.base.event_shape_tensor()
+
+  def _event_shape(self):
+    return self.base.event_shape
+
   def _quantization_offset(self):
     # Picks the "peakiest" of the component quantization offsets.
     offsets = helpers.quantization_offset(self.components_distribution)
@@ -225,6 +242,11 @@ def _lower_tail(self, tail_mass):
   def _upper_tail(self, tail_mass):
     return helpers.upper_tail(self.base, tail_mass)
 
+  @classmethod
+  def _parameter_properties(cls, dtype=tf.float32, num_classes=None):
+    raise NotImplementedError(
+        f"`{cls.__name__}` does not implement `_parameter_properties`.")
+
 
 class NoisyNormal(UniformNoiseAdapter):
   """Gaussian distribution with additive i.i.d. uniform noise."""