Skip to content

fix: ensure bounds vmaps as best as possible, more equinox errors, types for __eq__, _BoundsD/I#250

Open
beckermr wants to merge 77 commits into
mainfrom
bounds-fix-again-all-static
Open

fix: ensure bounds vmaps as best as possible, more equinox errors, types for __eq__, _BoundsD/I#250
beckermr wants to merge 77 commits into
mainfrom
bounds-fix-again-all-static

Conversation

@beckermr
Copy link
Copy Markdown
Collaborator

@beckermr beckermr commented May 20, 2026

This PR fixes a few other things with the bounds classes, images and typing.

  • ensure BoundsD/I vmaps as best as possible
  • use equinox to raise errors in the image classes
  • ensure all __eq__ methods return JAX boolean values except when no values from the class are ever traced (e.g., two static BoundsI classes, interpolants)
  • docs
  • implement _BoundsD/I
  • add tests
    • bool __eq__ API
    • bounds vmap
    • bounds set
    • bounds includes

closes #245
closes #237

@codspeed-hq
Copy link
Copy Markdown

codspeed-hq Bot commented May 20, 2026

Merging this PR will degrade performance by 34.94%

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

❌ 1 (👁 1) regressed benchmark
✅ 35 untouched benchmarks

Performance Changes

Mode Benchmark BASE HEAD Efficiency
👁 Simulation test_benchmark_invert_ab_noraise[run] 523 µs 803.8 µs -34.94%

Comparing bounds-fix-again-all-static (afe868f) with main (e2b1029)

Open in CodSpeed

@beckermr beckermr changed the base branch from main to typing-inits-cleanup May 20, 2026 18:46
Comment thread jax_galsim/bounds.py Outdated
Comment thread tests/jax/test_api.py Outdated
Comment thread jax_galsim/bounds.py Outdated
Comment thread jax_galsim/bounds.py Outdated
Comment thread jax_galsim/bounds.py Outdated
Comment thread docs/sharp-bits.rst Outdated
Comment thread docs/sharp-bits.rst Outdated
Comment thread tests/jax/test_api.py Outdated
Comment thread tests/jax/test_api.py
@beckermr beckermr changed the title fix: ensure bounds vmaps as best as possible, more equinox errors, types for __eq__ fix: ensure bounds vmaps as best as possible, more equinox errors, types for __eq__, _BoundsD/I May 23, 2026
@beckermr beckermr marked this pull request as ready for review May 23, 2026 12:38
@beckermr beckermr requested a review from ismael-mendoza May 23, 2026 12:38
@beckermr
Copy link
Copy Markdown
Collaborator Author

OK @ismael-mendoza this one is ready for review. I hope this is the last major bit of work on these classes for a long time!

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.

implement _BoundsI and _BoundsD better runtime checking of inits

1 participant