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

Job	Run time
Unittests (Python 3.11, JAX 0.6.0)	25s
Unittests (Python 3.12, JAX 0.6.2)	27s
Unittests (Python 3.11, JAX 0.5.3)	32s
Unittests (Python 3.11, JAX nightly)	1m 59s
Unittests (Python 3.12, JAX 0.7.1)	30s
Unittests (Python 3.13, JAX 0.5.3)	34s
Unittests (Python 3.11, JAX 0.7.1)	33s
Unittests (Python 3.10, JAX 0.6.0)	32s
Unittests (Python 3.13, JAX 0.6.0)	27s
Unittests (Python 3.12, JAX 0.6.1)	29s
Unittests (Python 3.13, JAX 0.7.2)	1m 39s
Unittests (Python 3.13, JAX nightly)	2m 1s
Unittests (Python 3.10, JAX 0.6.2)	25s
Unittests (Python 3.11, JAX 0.6.1)	32s
Unittests (Python 3.11, JAX 0.7.0)	31s
Unittests (Python 3.12, JAX 0.7.2)	33s
Unittests (Python 3.10, JAX 0.6.1)	29s
Unittests (Python 3.12, JAX nightly)	1m 46s
Unittests (Python 3.13, JAX 0.7.0)	30s
Unittests (Python 3.13, JAX 0.6.1)	32s
Unittests (Python 3.11, JAX 0.7.2)	28s
Unittests (Python 3.12, JAX 0.7.0)	28s
Unittests (Python 3.12, JAX 0.5.3)	26s
Unittests (Python 3.13, JAX 0.6.2)	32s
Unittests (Python 3.11, JAX 0.6.2)	31s
Unittests (Python 3.13, JAX 0.7.1)	29s
Unittests (Python 3.12, JAX 0.6.0)	33s
Unittests (Python 3.10, JAX 0.5.3)	33s
	19m 26s

Provide feedback