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

RotatingFrame.vectorized_map_into_frame is not JAX compatible and has ambiguous name #364

Open
DanPuzzuoli opened this issue Oct 21, 2024 · 1 comment
Labels
bug Something isn't working

Comments

@DanPuzzuoli
Copy link
Collaborator

The problem

This method is not JAX compatible (simple fix).

It is also "ambiguous" in the sense that "vectorized_map" does not indicate the nature of the transformation - a vectorized map in this context could either be the state of a differential equation or an operator that acts on states, and both are subject to different transformation rules. It currently implements the operator transformation, and there is no way to do a vectorized state transformation.

Suggestion

JAX fix is easy - I think there is a single usage of np that can be changed to unp. Along with this though, tests need to be added that would have caught this. I don't believe this method is being tested anywhere.

For the ambiguity, this method should be renamed vectorized_map -> vectorized_operator, and the equivalent vectorized_state method should also be added. There also isn't an out_of_frame version of this function, so one should be added for each of the operator and state methods.

@DanPuzzuoli DanPuzzuoli added the bug Something isn't working label Oct 21, 2024
@DanPuzzuoli
Copy link
Collaborator Author

One thing to note here is that the "vectorized" transformations can be implemented by the standard RotatingFrame methods but with the frame operator being the vectorized version of the map X -> [F, X]. So, an elegant solution to implementing these methods may be to just internally construct a RotatingFrame instance using the vectorized frame, and then call the ordinary methods.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

No branches or pull requests

1 participant