Implement StickBreakingTransform.log_abs_det_jacobian()

[fritzo]

This is needed to do HMC over simplex-constrained random variables, e.g. for Dirichlet random variables.

It would really help to have tests #110 complete to test correctness of this tricky bit of math.

[fritzo]

This is being implemented in pytorch#5547

[fritzo]

This was implemented in pytorch#5547 but may be buggy and does not work. We should add tests and fix the bugs.

@alshedivat Have you used StickBreakingTransform.log_abs_det_jacobian() at all? Do you have any ideas for tests?

EDIT sorry, I'm not sure where my bugs are, and they may not be in this function. But I would feel more confident if we had some tests for this function 🙂

[alshedivat]

@fritzo, my implementation is based on Stan's reference manual, v2.17.0 (Section 35.6, pp. 408-410) and uses the inverted stick-breaking construction. I also used an offset as suggested in the manual to make sure all zeros vector maps to the center of the simplex. I have used LogisticNormal distribution which is based on this transform.

Note: The original implementation of the StickBreakingTransform (before pytorch#5547) was slightly different. Could the bugs be due to assuming a particular form of the transform? (E.g., assuming no offset?). Could you elaborate more on what the bug is?

Edward uses the same construction in their to_simplex util function. We can borrow some tests from there, but I'm not sure if they have log_abs_det_jacobian implemented and tested.

@dustinvtran, do you happen to know a good way to test StickBreakingTransform? How is it done in Edward or Stan? (I can't find the relevant code.)

[fritzo]

@alshedivat Thanks for explaining. @ssnl pointed out that we could use torch.slogdet, but we would need to deal with the Jacobian not being square.

[dustinvtran]

Thanks for pinging me. We haven't implemented and tested a log_abs_det_jacobian for the stick breaking transform. I wholeheartedly agree with this direction though :) Stick breaking transforms are great.