ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

Module: tfp.experimental.util

TensorFlow Probability experimental python utilities.


class DeferredModule: Wrapper to defer initialization of a tf.Module instance.

class JitPublicMethods: Wrapper to compile an object's public methods using XLA.


as_composite(...): Returns a CompositeTensor equivalent to the given object.

make_trainable(...): Constructs a distribution or bijector instance with trainable parameters.

register_composite(...): A decorator that registers a TFP object as composite-friendly.

DEFAULT_METHODS_EXCLUDED_FROM_JIT ('event_shape', 'event_shape_tensor', 'batch_shape', 'batch_shape_tensor', 'dtype', 'kl_divergence', 'experimental_default_event_space_bijector', 'forward_event_shape', 'forward_event_shape_tensor', 'inverse_event_shape', 'inverse_event_shape_tensor', 'forward_dtype', 'inverse_dtype', 'forward_event_ndims', 'inverse_event_ndims')