Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Further signal arraylias fix #3

Merged

Conversation

DanPuzzuoli
Copy link

Summary

First, a failing test for a perturbative solvers using JAX has been fixed by removing Array from the envelope function and instead using unp.

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 calling unp.append(x, y) with x being a numpy array and y being a JAX array. Due to unp dispatching on the first argument, this results in the evaluation of np.append(x, y), which fails when y 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 into dynamics/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.

@DanPuzzuoli
Copy link
Author

Okay I think I got the docs to work as well. Sorry I think now this branch has conflicts with yours, though I suspect we made the same changes in the files with conflicts.

@to24toro to24toro merged commit 2ca2866 into to24toro:arraylias/signal Oct 27, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants