tfp.experimental.register_composite

View source on GitHub

A decorator that registers a Distribution as composite-friendly.

tfp.experimental.register_composite()

This registration is not required to call as_composite on instances of a given distribution, but it is required if a SavedModel with functions accepting or returning composite wrappers of this distribution will be loaded in python (without having called as_composite already).

Example:

class MyDistribution(tfp.distributions.Distribution):
   ...

# This will fail to load.
model = tf.saved_model.load(
    '/path/to/sm_with_funcs_returning_composite_tensor_MyDistribution')

Instead:

@tfp.experimental.register_composite
class MyDistribution(tfp.distributions.Distribution):
   ...

# This will load.
model = tf.saved_model.load(
    '/path/to/sm_with_funcs_returning_composite_tensor_MyDistribution')

Args
cls	A subclass of Distribution.

Returns
The input, with the side-effect of registering it as a composite-friendly distribution.

Raises
TypeError	If cls is not a subclass of Distribution, or if registration fails (cls is not convertible).
NotImplementedError	If registration fails (cls is not convertible).