Add copy method to Bijectors.

This duplicates the interface and behavior of Distribution.copy(). The implementation was copied from Distribution, then slightly simplified because Bijectors don't currently support batch slicing.

PiperOrigin-RevId: 380024162

Author: davmre

tensorflow_probability/python/bijectors/bijector.py

@@ -943,6 +943,28 @@ def __call__(self, value, name=None, **kwargs):
 
     return self.forward(value, name=name or 'forward', **kwargs)
 
+  def copy(self, **override_parameters_kwargs):
+    """Creates a copy of the bijector.
+
+    Note: the copy bijector may continue to depend on the original
+    initialization arguments.
+
+    Args:
+      **override_parameters_kwargs: String/value dictionary of initialization
+        arguments to override with new values.
+
+    Returns:
+      bijector: A new instance of `type(self)` initialized from the union
+        of self.parameters and override_parameters_kwargs, i.e.,
+        `dict(self.parameters, **override_parameters_kwargs)`.
+    """
+    parameters = dict(self.parameters, **override_parameters_kwargs)
+    b = type(self)(**parameters)
+    # pylint: disable=protected-access
+    b._parameters = self._no_dependency(parameters)
+    # pylint: enable=protected-access
+    return b
+
   def _forward_event_shape_tensor(self, input_shape):
     """Subclass implementation for `forward_event_shape_tensor` function."""
     # By default, we assume event_shape is unchanged.

tensorflow_probability/python/bijectors/bijector_test.py

@@ -183,6 +183,28 @@ def _parameter_properties(cls, dtype, num_classes=None):
     self.assertAllClose(ildj, true_ildj)
     self.assertAllClose(ildj, -fldj)
 
+  def testCopyExtraArgs(self):
+    # Note: we cannot easily test all bijectors since each requires
+    # different initialization arguments. We therefore spot test a few.
+    sigmoid = tfb.Sigmoid(low=-1., high=2., validate_args=True)
+    self.assertEqual(sigmoid.parameters, sigmoid.copy().parameters)
+    chain = tfb.Chain(
+        [
+            tfb.Softplus(hinge_softness=[1., 2.], validate_args=True),
+            tfb.MatrixInverseTriL(validate_args=True)
+        ], validate_args=True)
+    self.assertEqual(chain.parameters, chain.copy().parameters)
+
+  def testCopyOverride(self):
+    sigmoid = tfb.Sigmoid(low=-1., high=2., validate_args=True)
+    self.assertEqual(sigmoid.parameters, sigmoid.copy().parameters)
+    unused_sigmoid_copy = sigmoid.copy(validate_args=False)
+    base_params = sigmoid.parameters.copy()
+    copy_params = sigmoid.copy(validate_args=False).parameters.copy()
+    self.assertNotEqual(
+        base_params.pop('validate_args'), copy_params.pop('validate_args'))
+    self.assertEqual(base_params, copy_params)
+
 
 class IntentionallyMissingError(Exception):
   pass