Skip to content

Feature/jax speed up general#182

Merged
Jammy2211 merged 5 commits intofeature/jax_wrapperfrom
feature/jax_speed_up_general
Jul 12, 2025
Merged

Feature/jax speed up general#182
Jammy2211 merged 5 commits intofeature/jax_wrapperfrom
feature/jax_speed_up_general

Conversation

@Jammy2211
Copy link
Owner

Performance Improvements for General Use and Unit Testing

  • Significantly reduced runtime by avoiding unnecessary conversions between NumPy and JAX arrays within internal functions.
  • Updated several utility functions to dynamically select between numpy and jax.numpy based on the input array type. This is a temporary solution until the full JAX transition is complete.
  • Moved costly imports (e.g. from astropy) inside functions to defer and potentially avoid them during module import.
  • These changes noticeably improve both general runtime performance and unit test speed.

Further optimizations will follow once the final JAX migration is complete.

@Jammy2211 Jammy2211 requested a review from rhayes777 June 25, 2025 16:38
@Jammy2211 Jammy2211 merged commit 9a316ef into feature/jax_wrapper Jul 12, 2025
0 of 8 checks passed
@Jammy2211 Jammy2211 deleted the feature/jax_speed_up_general 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