hxtorch will soon depend on jax

We have been extending hxtorch.snn towards supporting event-based processing; we use jax for this purpose.

This MR adds this dependency of hxtorch on py-jax and py-tree-math.