Fixed a couple bugs with `pathwaysutils.jax` #348

Job	Run time
Unittests (Python 3.10, JAX 0.6.2)	32s
Unittests (Python 3.11, JAX 0.7.0)	28s
Unittests (Python 3.11, JAX 0.7.2)	24s
Unittests (Python 3.11, JAX nightly)	49s
Unittests (Python 3.12, JAX 0.5.3)	40s
Unittests (Python 3.10, JAX 0.6.1)	37s
Unittests (Python 3.13, JAX 0.7.0)	35s
Unittests (Python 3.12, JAX 0.6.1)	34s
Unittests (Python 3.12, JAX 0.6.0)	32s
Unittests (Python 3.12, JAX 0.6.2)	30s
Unittests (Python 3.11, JAX 0.5.3)	31s
Unittests (Python 3.10, JAX 0.5.3)	29s
Unittests (Python 3.12, JAX 0.7.2)	27s
Unittests (Python 3.13, JAX 0.6.2)	33s
Unittests (Python 3.13, JAX 0.6.1)	34s
Unittests (Python 3.10, JAX 0.6.0)	34s
Unittests (Python 3.11, JAX 0.7.1)	24s
Unittests (Python 3.13, JAX 0.5.3)	34s
Unittests (Python 3.11, JAX 0.6.0)	31s
Unittests (Python 3.13, JAX 0.6.0)	27s
Unittests (Python 3.12, JAX 0.7.1)	26s
Unittests (Python 3.13, JAX 0.7.1)	30s
Unittests (Python 3.12, JAX 0.7.0)	35s
Unittests (Python 3.13, JAX nightly)	2m 1s
Unittests (Python 3.12, JAX nightly)	1m 58s
Unittests (Python 3.11, JAX 0.6.1)	32s
Unittests (Python 3.11, JAX 0.6.2)	30s
Unittests (Python 3.13, JAX 0.7.2)	34s
	17m 51s

Provide feedback