Skip to content

Feature/xp no autofit import#193

Merged
Jammy2211 merged 4 commits intomainfrom
feature/xp_no_autofit_import
Nov 13, 2025
Merged

Feature/xp no autofit import#193
Jammy2211 merged 4 commits intomainfrom
feature/xp_no_autofit_import

Conversation

@Jammy2211
Copy link
Owner

A recent PR on the child projects made JAX optional for likelihood functions, whereby users pass the JAX namespace as the variable xp through the source code.

This PR makes JAX optional at the highest level (e.g. PyAutoConf and PyAutoFit), including:

  • For a non-linear search to use JAX, the use_jax input must be passed as True to the Analysis object.
  • The non-linear search will internally work out if it supports JAX natively. This will ultimately have behavior where, for example, if gradients are used it uses jax.grad, if not it uses jax.jit, and if batching is support jax.vmap.
  • Currently only Nautilus uses the Analysis.use_jax attribute to set up a jax.vmap.

There are few hacky unclean bits in the autofit model composition where it determines whether to use JAX based on input type. A more thorough consideration of how JAX should work in autofit will be performed in the future.

@Jammy2211 Jammy2211 merged commit a5e3e78 into main Nov 13, 2025
4 of 8 checks passed
@Jammy2211 Jammy2211 deleted the feature/xp_no_autofit_import branch November 30, 2025 16:31
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