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

[leventov]

Describe the issue:

#133 In this PR, the second branch of this conditional:

pytensor/tests/link/jax/test_basic.py

Lines 73 to 77 in d175203

	if must_be_device_array:
	if isinstance(jax_res, list):
	assert all(isinstance(res, jax.Array) for res in jax_res)
	else:
	assert isinstance(jax_res, jax.interpreters.xla.DeviceArray)

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.