I think it should just be:
jac = tg.jacobian(f_inv_x, x) # (10, 2, 2)
_, logdet = pt.linalg.slogdet(jac) # (10, )
return logdet
I think it should just be:
jac = tg.jacobian(f_inv_x, x) # (10, 2, 2)
_, logdet = pt.linalg.slogdet(jac) # (10, )
return logdet