fix jax and numpy ambiguity in tree builder

Currently we are using both jnp and np for arrays. This has a known effect on precision, and could be the cause of some inconsistencies in the patient traversal.

shows up [here](https://github.com/DoYouEvenStackSmash/data-analysis/blob/ba4ce2844fbd70a75b580e6258fef43793d1daaa/src/tree-likelihood/python/tree_build.py#L81)