Skip to content

Feature/jax messages#1165

Open
Jammy2211 wants to merge 15 commits intomainfrom
feature/jax_messages
Open

Feature/jax messages#1165
Jammy2211 wants to merge 15 commits intomainfrom
feature/jax_messages

Conversation

@Jammy2211
Copy link
Collaborator

Makes it so that the messages package supports JAX via the xp interface implemented elsewhere in PyAuto repos.

This PR went OK, considering how far through the depths of the graphical source code the changes went. Integration tests on all this functionality seem to perform correctly but now work with JAX.

Its worth noting an environment variable implementation would not work, because for a given EP or graphical run this code is cycled through using both numpy and JAX (only uses JAX inside a non linear search during modeling).

There are a few nasty hacks which I need to better understand JAX in order to clean up.

@Jammy2211 Jammy2211 requested a review from rhayes777 November 30, 2025 16:22
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.

1 participant