def get_transform(self, params):
def loss_fn(z):
params1 = params.copy()
params1['{}_loc'.format(self.prefix)] = z
return self._loss_fn(params1)
loc = params['{}_loc'.format(self.prefix)]
precision = hessian(loss_fn)(loc)
scale_tril = cholesky_of_inverse(precision)
if not_jax_tracer(scale_tril):
if np.any(np.isnan(scale_tril)):
warnings.warn("Hessian of log posterior at the MAP point is singular. Posterior"
" samples from AutoLaplaceApproxmiation will be constant (equal to"
" the MAP point).")
scale_tril = jnp.where(jnp.isnan(scale_tril), 0., scale_tril)
return LowerCholeskyAffine(loc, scale_tril)
def get_transform(self, params):
posterior = self.get_posterior(params)
return LowerCholeskyAffine(posterior.loc, posterior.scale_tril)
def get_transform(self, params):
loc = params["{}_loc".format(self.prefix)]
scale_tril = params["{}_scale_tril".format(self.prefix)]
return LowerCholeskyAffine(loc, scale_tril)