Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
First, a failing test for a perturbative solvers using JAX has been fixed by removing
Array
from the envelope function and instead usingunp
.The rest of this PR fixes the remaining failures that all stem from the same issue (in both the pulse -> signal conversion code and in
operator_collections.py
). These were all the result of callingunp.append(x, y)
withx
being anumpy
array andy
being a JAX array. Due tounp
dispatching on the first argument, this results in the evaluation ofnp.append(x, y)
, which fails wheny
is a JAX tracer.The issue here is that what we fundamentally want here is multiple dispatch, but arraylias is single dispatch. To resolve this, I've added functions
_preferred_lib
and_numpy_multi_dispatch
intodynamics/arraylias/alias
that, given a list of arrays as arguments, will determine which array backend to use. Basically, for now it works as: if all arrays are of type"numpy"
use numpy, but if any are of type"jax"
then JAX is used. The current logic presupposes that these are the only two libraries. This will change later if more libraries are added.