Skip to content

Feature/xp no autofit import#78

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

Feature/xp no autofit import#78
Jammy2211 merged 2 commits intomainfrom
feature/xp_no_autofit_import

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

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 requested a review from rhayes777 November 13, 2025 11:16
@codecov
Copy link
Copy Markdown

codecov bot commented Nov 13, 2025

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.81%. Comparing base (afca90a) to head (a55fea4).
⚠️ Report is 14 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main      #78      +/-   ##
==========================================
- Coverage   79.79%   77.81%   -1.98%     
==========================================
  Files          13       14       +1     
  Lines         787      807      +20     
==========================================
  Hits          628      628              
- Misses        159      179      +20     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@Jammy2211 Jammy2211 merged commit 53dfc99 into main Nov 13, 2025
9 of 10 checks passed
@Jammy2211 Jammy2211 deleted the feature/xp_no_autofit_import branch February 13, 2026 13:37
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