Skip to content

refactor: defer nufftax/jax and numba imports to first use - #451

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/lazy-heavy-imports
Aug 19, 2026
Merged

refactor: defer nufftax/jax and numba imports to first use#451
Jammy2211 merged 1 commit into
mainfrom
feature/lazy-heavy-imports

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Defers PyAutoArray's two heavy import-time costs to first use, as part of the cross-repo import-time task (PyAutoFit#1505): the nufftax import (which pulls jax, ~0.7s) and numba jit decoration (~0.2s). import autoarray drops to ~1.15s with neither jax, nufftax nor numba loading on bare import.

  • operators/transformer.py: the module-level import nufftax and the version-gated 0.6.x batcher shim move into an idempotent _load_nufftax(), called from TransformerNUFFT.__init__ and the three methods that dereference the module global (_forward_native, image_from, transform_mapping_matrix) — instances unpickled in multiprocessing workers never re-run __init__, so the method anchors are load-bearing.
  • numba_util.py: jit() queues functions and hands them to numba on the first call of any decorated function, rebinding each defining module's global to the real dispatcher so nopython cross-calls (5 exist) still resolve at compile time. ModuleNotFoundError → plain function fallback preserved; config-driven nopython/cache/parallel options unchanged.

API Changes

None — internal changes only. All 29 @numba_util.jit() functions and TransformerNUFFT behave identically; laziness only moves when the underlying imports happen.
See full details below.

Test Plan

  • Full suite: 1062 passed
  • TransformerNUFFT numpy/jax visibilities agree exactly; transform_mapping_matrix + image_from verified; pickled-instance-with-reset-module-state (worker simulation) works
  • numba cross-call path (binned_mask_frombinned_image_from) compiles and runs; module global rebinds to CPUDispatcher
  • Downstream autogalaxy (1111) and autolens (538) suites green against this branch
Full API Changes (for automation & release notes)

Changed Behaviour

  • autoarray.numba_util.jit — decorated functions are lightweight wrappers until first call (then real numba dispatchers, module globals rebound). Dispatcher-only attributes (e.g. .py_func) are unavailable before first call; no users exist in the organism.
  • autoarray.operators.transformernufftax (and jax via it) import on first TransformerNUFFT use instead of at module import; a missing nufftax now errors at construction (same nufftax_exception()), not silently at import.

Migration

  • None required.

Generated by the PyAutoLabs agent workflow.

…#1505)

_load_nufftax() anchors the nufftax import + 0.6.x batcher shim at
TransformerNUFFT use (init + the three module-global-dereferencing methods,
covering unpickled worker instances); numba_util.jit lazily materializes
all decorated functions on first call, rebinding module globals so
nopython cross-calls resolve.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Aug 19, 2026
@Jammy2211
Jammy2211 merged commit a26385b into main Aug 19, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/lazy-heavy-imports branch August 19, 2026 22:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant