Remove tfp.experimental.as_composite, as it does nothing (has been st…

…ubbed out).

With few exceptions (classes with callbacks mostly), all TensorFlow Probability (tfp) distributions and TensorFLow linear operators are already composite tensors, so this conversion function is no longer necessary. For out-of-tfp distributions, previous changes have already made the manual changes needed to make them be composite tensors.

PiperOrigin-RevId: 545736125
Change-Id: Iecb2d02a58db91a9800153b493a4d41eb2189e7b

acme/tf/savers.py

@@ -27,7 +27,6 @@
 from acme.utils import paths
 import sonnet as snt
 import tensorflow as tf
-import tensorflow_probability as tfp
 import tree
 
 from tensorflow.python.saved_model import revived_types
@@ -375,19 +374,10 @@ def make_snapshot(module: snt.Module):
          'which is required for snapshotting; run '
          'create_variables to add this annotation.').format(module.name))
 
-  # This function will return the object as a composite tensor if it is a
-  # distribution and will otherwise return it with no changes.
-  def as_composite(obj):
-    if isinstance(obj, tfp.distributions.Distribution):
-      return tfp.experimental.as_composite(obj)
-    else:
-      return obj
-
-  # Replace any distributions returned by the module with composite tensors and
-  # wrap it up in tf.function so we can process it properly.
+  # Wrap the module up in tf.function so we can process it properly.
   @tf.function
   def wrapped_module(*args, **kwargs):
-    return tree.map_structure(as_composite, module(*args, **kwargs))
+    return module(*args, **kwargs)
 
   # pylint: disable=protected-access
   snapshot = Snapshot()