1e-16 is too much precision to ask for. You don’t get that precision on addition and multiplication. Probably cutting off at 1e-8 or so precision would be enough and a lot cheaper to calculate. And you probably want the precision to be proportional, not absolute. Something like when s stops chanign the way you have it here.
As the algorithm’s written, you could store the si and use a single log_sum_exp. That’s a lot less arithmetic.
Move the definition of log_z and test up to first two operations after defining s. That helps with clarity and will save some operations when the condition triggers.
Rather than all the k * log_z, you can just keep a running total and add log_z at a time.
None of this other than fixing precision will matter much compared to the cost of lgamma.