https://github.com/harvardnlp/pytorch-struct/blob/8c8c14e49360ff74c127ad8e64fa7fc3ae81ca73/torch_struct/distributions.py#L232 I think `lengths=None` should be just `lengths`?