HeteroskedasticTFPConditional should construct tensors at class-construction, not at module-import time
@@ -97,7 +97,7 @@ class HeteroskedasticTFPConditional(MultiLatentTFPConditional):
def __init__(
self,
distribution_class: Type[tfp.distributions.Distribution] = tfp.distributions.Normal,
- scale_transform: tfp.bijectors.Bijector = positive(base="exp"),
+ scale_transform: Optional[tfp.bijectors.Bijector] = None,
**kwargs,
):
"""
@@ -105,13 +105,16 @@ class HeteroskedasticTFPConditional(MultiLatentTFPConditional):
as first and second argument, respectively.
:param scale_transform: callable/bijector applied to the latent
function modelling the scale to ensure its positivity.
- Typically, `tf.exp` or `tf.softplus`, but can be any function f: R -> R^+.
+ Typically, `tf.exp` or `tf.softplus`, but can be any function f: R -> R^+. Defaults to exp if not explicitly specified.
"""
+ if scale_transform is None:
+ scale_transform = positive(base="exp")
+ self.scale_transform = scale_transform
def conditional_distribution(Fs) -> tfp.distributions.Distribution:
tf.debugging.assert_equal(tf.shape(Fs)[-1], 2)
loc = Fs[..., :1]
- scale = scale_transform(Fs[..., 1:])
+ scale = self.scale_transform(Fs[..., 1:])
return distribution_class(loc, scale)
super().__init__(
