Skip to content

feat: SimulatorInterferometer.via_tracer_from auto-default xp from parent use_jax#540

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/simulator-interferometer-use-jax
May 24, 2026
Merged

feat: SimulatorInterferometer.via_tracer_from auto-default xp from parent use_jax#540
Jammy2211 merged 1 commit into
mainfrom
feature/simulator-interferometer-use-jax

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

Companion PR to the PyAutoArray PR adding use_jax=True to aa.SimulatorInterferometer. The autolens via_tracer_from and via_galaxies_from overrides gain xp=None parameter and default to self._xp from the parent.

After both PRs merge:

simulator = al.SimulatorInterferometer(
    uv_wavelengths=uv_wavelengths, exposure_time=300.0, use_jax=True
)
dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)  # JAX-backed

API Changes

  • Changed signature: al.SimulatorInterferometer.via_tracer_from(tracer, grid, xp=None) and via_galaxies_from(galaxies, grid, xp=None) — both gain xp parameter; xp=None falls back to self._xp.

See full details below.

Test Plan

  • PyAutoLens 317/317 tests pass.
  • Eager cross-xp parity via the workspace_test parity script.
Full API Changes

Changed signature

  • al.SimulatorInterferometer.via_tracer_from(tracer, grid, xp=None).
  • al.SimulatorInterferometer.via_galaxies_from(galaxies, grid, xp=None).

Migration

  • No breaking changes. Existing callers continue to work.

Depends on

  • Jammy2211/PyAutoArray PR adding use_jax to aa.SimulatorInterferometer. Merge PyAutoArray PR first.

🤖 Generated with Claude Code

…arent

Changes via_tracer_from / via_galaxies_from on autolens's SimulatorInterferometer
override to default xp=None and fall back to self._xp inherited from
aa.SimulatorInterferometer. When the simulator is constructed with use_jax=True,
both methods route xp=jnp through tracer.image_2d_from / galaxies.image_2d_from
and via_image_from without the user passing xp explicitly.

Depends on Jammy2211/PyAutoArray PR adding use_jax constructor flag + _xp
property to the parent aa.SimulatorInterferometer.

Part of Phase 2 PR 3 of z_features/jax_user_intro.md.

Issue: PyAutoArray#334

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label May 24, 2026
@Jammy2211 Jammy2211 merged commit db7fdc9 into main May 24, 2026
4 of 6 checks passed
@Jammy2211 Jammy2211 deleted the feature/simulator-interferometer-use-jax branch May 24, 2026 16:00
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