minor bug: jax tests reference jax.interpreters.xla.DeviceArray

[leventov]

Describe the issue:

https://github.com/pymc-devs/pytensor/pull/133 In this PR, the second branch of this conditional:

https://github.com/pymc-devs/pytensor/blob/d175203b4e00f48db9c61b68a5f70263a1fbb645/tests/link/jax/test_basic.py#L73-L77

was evidently not updated only because it is never executed in the tests.

cc @ricardoV94

Reproducable code example:

n/a

Error message:

No response

PyTensor version information:

n/a

Context for the issue:

No response

[ricardoV94]

@leventov want to update it?

[leventov]

Yes, I can do it later.

Also: probably this function as well as compare_numba_and_py() should use itertools.zip_longest rather than zip.