Arraylias integration - Signal class -

[to24toro]

Summary

I started to integrate arraylias to dynamics.
This PR includes the integration of the Signal class in dynamics.

Details and comments

[to24toro]

This decorator works as creating TestDiscreteSignalSum of "numpy", "jax", "array_numpy", "array_jax", which inherits TestSignalSum. But there is no class whose name is TestSignalSum because it has already been decorated by test_array_backends.

[to24toro]

The current implementation fails the tutorial in optimizing_pulse_sequence.rst. The main reason for this failure is that numpy.array and jax.numpy.array sometimes gets in the mix in the code. This would be no problem if we don't use jax.jit. However, this failure is found within jit(value_and_grad(objective)) in this tutorial.

• In the existing tutorial code, dynamics.Array is used, but for example, even if dynamics.Array backend is jax, it will be converted to numpy.array by unp.asarray(). I think this is at https://github.com/Qiskit-Extensions/qiskit-dynamics/blob/3a373cfb834cd2f2e15245ab924aa3568b8db326/qiskit_dynamics/array/array.py#L188 .
  If we are to continue supporting dynamics.Array by the completion of arraylias integration, we will need to change or improve the dynamics.Array .

• Suppose we can change samples used in DIscreteSignal to jax.numpy.array through unp.asarray. In that case, an error will occur if there is a numpy.array elsewhere. Moreover, taking Convolution as an example, in the following code:

if isinstance(signal, DiscreteSignal):
  # Perform a discrete time convolution.
  dt = signal.dt
  func_samples = np.asarray([self._func(dt * i) for i in range(signal.duration)])
  func_samples = func_samples / sum(func_samples)
  sig_samples = signal(dt * np.arange(signal.duration))

even if np -> unp is changed, unp.arange(signal.duration) becomes a numpy.array because signal.duration is assumed to be an int. The point is that types like int or float get coverted to numpy by unp.arange(). In the other words, numpy takes precedence over jax, leading to the issue.

[DanPuzzuoli]

The signals and test/dynamics/signals folders look good to me aside from the two comments I made:

• change alias to numpy_alias
• Remove all redundant type hinting for "numbers" and List where ArrayLike appears.

[DanPuzzuoli]

I think maybe we should call this numpy_alias instead of just alias now that there will be one for NumPy and one for SciPy.

[to24toro]

fixed at 79c45ce .

[DanPuzzuoli]

I think the type hints are redundant as ArrayLike includes float, int, complex. Should we remove everything but ArrayLike here?

I see ArrayLike doesn't include List, but maybe we should add that, as it does technically act on lists as well.

Edit: Ah yes I see later on in the file for DiscreteSignalSum you do conversions like Union[float, List, ArrayLike] -> ArrayLike. This should be done for all similar type hints throughout the whole file.

[to24toro]

You are right. I fixed at a3509f5 and 921126b .

[DanPuzzuoli]

I'm trying to figure out how to get the tests to pass in other submodules - the issue with a lot of the failures seems to be that:

unp.asarray(Array(x, backend='jax'))

ends up calling numpy.array(x).

This issue comes up in tests like test.dynamics.models.test_generator_model.TestGeneratorModelSparseJax.test_jit_grad in which a value a is being traced, and the function contains the code:

Signal(Array(a))

This test (and I'm guessing many others) can be fixed by simply changing the above to:

Signal(a)

This is how we'll want to do things going forward anyway, but this doesn't help with the fact that Signal(Array(x, backend='jax')) is broken when x is a tracer.

Fixing this by modifying Array appears to be nontrivial. I haven't been able to find a quick fix as numpy gets upset when np.array returns something that is not an array. As such, it seems that maybe asarray cannot work how we'd want for Array with backend='jax'.

[DanPuzzuoli]

I've figured out how to bypass the above issue and directly fix most tests here. The only thing I haven't fixed yet is stuff to do with pulse -> signal conversion and JAX compatibility - but I think it may just be the same issue.

I'll attempt again to see if we can "properly" fix this problem (i.e. get unp.asarray to work with Array with backend=="jax"). It would be nice to have backwards compatibility with code that uses Array, but I'm also okay with it breaking if we can't find a clean solution.

[to24toro]

The tests which have not been passed are two types:

• When jitting, https://github.com/Qiskit-Extensions/qiskit-dynamics/blob/1e5fae945caac10b99e7c325e25bf38067971b9c/qiskit_dynamics/pulse/pulse_to_signals.py#L169 is the resource of errors because this list is converted to numpy.array. The solution is implementing setting something like Array.default_backend() and when it is "jax", we need to pass jnp.array([]) to samples.
• np.append(ham_sig_vals, dis_sig_vals, axis=-1) in "qiskit_dynamics/models/operator_collections.py". Should it be solved when working on operator_collections?

[DanPuzzuoli]

All looking good to me, but some final nitpicking details to sort out.

[DanPuzzuoli]

I suspect these lines may be prone to the dispatching errors we had to deal with in other places. Should these be changed to use _numpy_multi_dispatch?

[to24toro]

I forgot to change it to _numpy_multi_dispatch.
I fixed at 2e1436e.

[DanPuzzuoli]

_numpy_multi_dispatch?

[to24toro]

I fixed at 2e1436e.

[DanPuzzuoli]

_numpy_multi_dispatch?

[to24toro]

I fixed at 2e1436e.

[DanPuzzuoli]

_numpy_multi_dispatch?

[to24toro]

I fixed at 2e1436e.