def _piecewise_cdf(self, inputs, transform_params, inverse=False):
unnormalized_pdf = transform_params
if self.tails is None:
return splines.linear_spline(inputs=inputs,
unnormalized_pdf=unnormalized_pdf,
inverse=inverse)
else:
return splines.unconstrained_linear_spline(
inputs=inputs,
unnormalized_pdf=unnormalized_pdf,
inverse=inverse,
tails=self.tails,
tail_bound=self.tail_bound,
)
def _spline(self, inputs, inverse=False):
batch_size = inputs.shape[0]
unnormalized_pdf = _share_across_batch(self.unnormalized_pdf, batch_size)
if self.tails is None:
outputs, logabsdet = splines.linear_spline(
inputs=inputs, unnormalized_pdf=unnormalized_pdf, inverse=inverse
)
else:
outputs, logabsdet = splines.unconstrained_linear_spline(
inputs=inputs,
unnormalized_pdf=unnormalized_pdf,
inverse=inverse,
tails=self.tails,
tail_bound=self.tail_bound,
)
return outputs, torchutils.sum_except_batch(logabsdet)
def call_spline_fn(inputs, inverse=False):
return splines.linear_spline(inputs=inputs,
unnormalized_pdf=unnormalized_pdf,
inverse=inverse)
