Replace deprecation of parameter_properties inheritance with a warning.

This also removes a spurious deprecation warning that appears when `JointDistribution`s are used in public colabs. (JDs have a multi-level hierarchy in which no class defines its own _parameter_properties, but since they all inherit the base class def that raises `NotImplementedError`, there's no inheritance problem).

Justification for this change: there are legitimate 'quick-and-dirty' uses of parameter_properties inheritance, even if we wouldn't do it in TFP.

For example, somewhere in the DeepMind silo is a Normal distribution that takes a log_scale rather than a scale:

class NormalWithLogScale(tfd.Normal):

  def __init__(self, loc, log_scale):
    super().__init__(loc=loc, scale=tf.exp(log_scale))

Ignoring whether this is the best way to accomplish any particular goal, from a general Pythonic standpoint one might expect that this class would at least be basically functional. It may not support batch slicing or AutoCompositeTensor, but one should at least be able to call sample and log_prob, which means it should at least define properties like batch_shape. But now that batch shape depends on parameter_properties (as of cl/373590501), breaking parameter_properties inheritance means breaking batch_shape inheritance.

To allow quick subclasses like this, I propose we simply warn when an inherited `parameter_properties` is called. In this example, the batch shape would be computed (correctly) using the base Normal parameters, just as if an explicit batch_shape method had been inherited.

For full functionality including batch slicing, CompositeTensor, etc., a subclass would need to both (a) set self.parameters = dict(locals()) in its own constructor, and (b) define its own _parameter_properties. I've tried to articulate these requirements in the warning message.

PiperOrigin-RevId: 374554087

Author: davmre

tensorflow_probability/python/distributions/distribution.py

@@ -23,6 +23,7 @@
 import contextlib
 import functools
 import inspect
+import logging
 import types
 
 import decorator
@@ -275,40 +276,61 @@ def __new__(mcs, classname, baseclasses, attrs):
       return super(_DistributionMeta, mcs).__new__(
           mcs, classname, baseclasses, attrs)
 
-    # Subclasses shouldn't inherit their parents' `_parameter_properties`,
-    # since (in general) they'll have different parameters. Exceptions (for
-    # convenience) are:
+    # Warn when a subclass inherits `_parameter_properties` from its parent
+    # (this is unsafe, since the subclass will in general have different
+    # parameters). Exceptions are:
     #  - Subclasses that don't define their own `__init__` (handled above by
     #    the short-circuit when `default_init is None`).
     #  - Subclasses that define a passthrough `__init__(self, *args, **kwargs)`.
-    #  - Direct children of `Distribution`, since the inherited method just
-    #    raises a NotImplementedError.
+    # pylint: disable=protected-access
     init_argspec = tf_inspect.getfullargspec(default_init)
     if ('_parameter_properties' not in attrs
-        and base != Distribution
         # Passthrough exception: may only take `self` and at least one of
         # `*args` and `**kwargs`.
         and (len(init_argspec.args) > 1
              or not (init_argspec.varargs or init_argspec.varkw))):
-      # TODO(b/183457779) remove warning and raise `NotImplementedError`.
-      attrs['_parameter_properties'] = deprecation.deprecated(
-          date='2021-07-01',
-          instructions="""
-Calling `_parameter_properties` on subclass {classname} that redefines the
-parent ({basename}) `__init__` is unsafe and will raise an error in the future.
-Please implement an explicit `_parameter_properties` for the subclass. If the
-subclass `__init__` takes the same parameters as the parent, you may use the
-placeholder implementation:
 
-  @classmethod
-  def _parameter_properties(cls, dtype, num_classes=None):
-    return {basename}._parameter_properties(
-        dtype=dtype, num_classes=num_classes)
+      @functools.wraps(base._parameter_properties)
+      def wrapped_properties(*args, **kwargs):  # pylint: disable=missing-docstring
+        """Wrapper to warn if `parameter_properties` is inherited."""
+        properties = base._parameter_properties(*args, **kwargs)
+        # Warn *after* calling the base method, so that we don't bother warning
+        # if it just raised NotImplementedError anyway.
+        logging.warning("""
+Distribution subclass %s inherits `_parameter_properties from its parent (%s)
+while also redefining `__init__`. The inherited annotations cover the following
+parameters: %s. It is likely that these do not match the subclass parameters.
+This may lead to errors when computing batch shapes, slicing into batch
+dimensions, calling `.copy()`, flattening the distribution as a CompositeTensor
+(e.g., when it is passed or returned from a `tf.function`), and possibly other
+cases. The recommended pattern for distribution subclasses is to define a new
+`_parameter_properties` method with the subclass parameters, and to store the
+corresponding parameter values as `self._parameters` in `__init__`, after
+calling the superclass constructor:
+
+```
+class MySubclass(tfd.SomeDistribution):
+
+  def __init__(self, param_a, param_b):
+    parameters = dict(locals())
+    # ... do subclass initialization ...
+    super(MySubclass, self).__init__(**base_class_params)
+    # Ensure that the subclass (not base class) parameters are stored.
+    self._parameters = parameters
+
+  def _parameter_properties(self, dtype, num_classes=None):
+    return dict(
+      # Annotations may optionally specify properties, such as `event_ndims`,
+      # `default_constraining_bijector_fn`, `specifies_shape`, etc.; see
+      # the `ParameterProperties` documentation for details.
+      param_a=tfp.util.ParameterProperties(),
+      param_b=tfp.util.ParameterProperties())
+```
+""", classname, base.__name__, str(properties.keys()))
+        return properties
 
-""".format(classname=classname,
-           basename=base.__name__))(base._parameter_properties)
+      attrs['_parameter_properties'] = wrapped_properties
 
-    # pylint: disable=protected-access
     # For a comparison of different methods for wrapping functions, see:
     # https://hynek.me/articles/decorators/
     @decorator.decorator
@@ -656,10 +678,7 @@ def _composite_tensor_shape_params(self):
   @classmethod
   def _parameter_properties(cls, dtype, num_classes=None):
     raise NotImplementedError(
-        '_parameter_properties` is not implemented: {}. '
-        'Note that subclasses that redefine `__init__` are not assumed to '
-        'share parameters with their parent class and must provide a separate '
-        'implementation.'.format(cls.__name__))
+        '_parameter_properties` is not implemented: {}.'.format(cls.__name__))
 
   @classmethod
   def parameter_properties(cls, dtype=tf.float32, num_classes=None):

tensorflow_probability/python/distributions/distribution_test.py

@@ -19,6 +19,7 @@
 import collections
 # Dependency imports
 
+from absl import logging
 from absl.testing import parameterized
 
 import numpy as np
@@ -29,8 +30,6 @@
 from tensorflow_probability.python.internal import test_util
 
 from tensorflow.python.framework import test_util as tf_test_util  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
-from tensorflow.python.platform import test as tf_test  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
-from tensorflow.python.platform import tf_logging  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
 
 
 class TupleDistribution(tfd.Distribution):
@@ -617,10 +616,7 @@ def normal_differential_entropy(scale):
                     self.evaluate(normal_differential_entropy(scale)),
                     err=1e-5)
 
-  @test_util.jax_disable_test_missing_functionality('tf_logging')
-  @tf_test.mock.patch.object(tf_logging, 'warning', autospec=True)
-  def testParameterPropertiesNotInherited(self, mock_warning):
-    # TODO(b/183457779) Test for NotImplementedError (rather than just warning).
+  def testParameterPropertiesNotInherited(self):
 
     # Subclasses that don't redefine __init__ can inherit properties.
     class NormalTrivialSubclass(tfd.Normal):
@@ -640,20 +636,27 @@ class MyDistribution(tfd.Distribution):
       def __init__(self, param1, param2):
         pass
 
-    NormalTrivialSubclass.parameter_properties()
-    NormalWithPassThroughInit.parameter_properties()
-    with self.assertRaises(NotImplementedError):
-      MyDistribution.parameter_properties()
-    self.assertEqual(0, mock_warning.call_count)
+    with self.assertLogs(level=logging.WARNING) as log:
+      NormalTrivialSubclass.parameter_properties()
+      NormalWithPassThroughInit.parameter_properties()
+      with self.assertRaises(NotImplementedError):
+        MyDistribution.parameter_properties()
+      with self.assertRaises(NotImplementedError):
+        # Ensure that the unimplemented JD propertoes don't raise a warning.
+        tfd.JointDistributionCoroutine.parameter_properties()
+      logging.warning('assertLogs context requires at least one warning.')
+    # Assert that no warnings occurred other than the dummy warning.
+    self.assertLen(log.records, 1)
 
     class NormalWithExtraParam(tfd.Normal):
 
       def __init__(self, extra_param, *args, **kwargs):
         self._extra_param = extra_param
         super(NormalWithExtraParam, self).__init__(*args, **kwargs)
 
-    NormalWithExtraParam.parameter_properties()
-    self.assertEqual(1, mock_warning.call_count)
+    with self.assertLogs(level=logging.WARNING) as log:
+      NormalWithExtraParam.parameter_properties()
+    self.assertLen(log.records, 1)
 
 
 @test_util.test_all_tf_execution_regimes