Dev#303
Conversation
…301) `BatchLM` uses `vmap(jacfwd(...))` to compute batched Jacobians. PyTorch's forward-mode AD cannot trace through in-place tensor mutations, causing failures when any `integrate_mode` other than `"none"` was used. ## Core fix: backend index operations `_fill_at_indices_torch` and `_add_at_indices_torch` now use `torch.index_put` (functional, out-of-place) for LongTensor indices — the path taken by `topk`-based adaptive integration. Bool/tuple/slice indices fall back to `clone()` + in-place, which preserves compatibility for non-differentiated callers. ```python # Before — breaks vmap+jacfwd def _fill_at_indices_torch(self, array, indices, values): array[indices] = values return array # After — compatible with vmap+jacfwd for LongTensor indices def _fill_at_indices_torch(self, array, indices, values): if isinstance(indices, self.module.Tensor) and indices.dtype != self.module.bool: return self.module.index_put(array, (indices,), values) result = array.clone() result[indices] = values return result ``` ## Sweep for bool-mask in-place operations in brightness functions `vmap` additionally rejects operations that produce dynamic shapes (bool-indexed selection). Fixed in three places: - **`spline.py`** — extrapolation boundary: replaced `fill_at_indices(I, R > profR[-1], 0)` with `backend.where(R > profR[-1], zeros, I)`. - **`WedgeMixin.polar_model`** — replaced `add_at_indices(model, bool_mask, iradial_model(s, R[mask]))` with `model + where(mask, iradial_model(s, R), zeros)`. The radial model is now evaluated over the full pixel array; the mask is applied via `where` to maintain static shapes. - **`RayMixin.polar_model`** — same pattern for cosine-weighted segment accumulation. ## New tests Added `BatchLM`-specific tests in `test_fit.py` covering: - All `integrate_mode` values: `"none"`, `"bright"`, `"curvature"` - Common `sampling_mode` values: `"midpoint"`, `"simpsons"`, `"quad:3"` - Images with rotated CD matrices (non-trivial WCS orientations) - Poisson likelihood path --------- Co-authored-by: Connor Stone, PhD <connorstone628@gmail.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: ConnorStoneAstro <78555321+ConnorStoneAstro@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR updates model evaluation to be compatible with BatchLM’s vmap+jacfwd execution (static shapes) and adds regression tests covering key BatchLM configurations.
Changes:
- Added BatchLM regression tests across
integrate_mode,sampling_mode, rotated CD matrices, Poisson likelihood, and basic loss improvement expectations. - Replaced index-based updates (
add_at_indices/fill_at_indices) with shape-stablewhere-based expressions in brightness mixins and spline evaluation. - Updated Torch/JAX backend index update helpers and removed
and_at_indiceswiring.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 9 comments.
| File | Description |
|---|---|
| tests/test_fit.py | Adds BatchLM regression tests and a helper to generate batch fitting setups. |
| astrophot/models/mixins/brightness.py | Switches from indexed updates to where/broadcasting to keep tensor shapes static under vmap+jacfwd. |
| astrophot/models/func/spline.py | Replaces in-place fill-at-indices with functional where for shape stability. |
| astrophot/backend_obj.py | Adjusts torch index update implementation for vmap+jacfwd compatibility and removes and_at_indices. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #303 +/- ##
==========================================
+ Coverage 91.23% 91.47% +0.23%
==========================================
Files 113 113
Lines 6198 6194 -4
==========================================
+ Hits 5655 5666 +11
+ Misses 543 528 -15
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
…304) Bumps [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) from 1.13.0 to 1.14.0. <details> <summary>Release notes</summary> <p><em>Sourced from <a href="https://github.com/pypa/gh-action-pypi-publish/releases">pypa/gh-action-pypi-publish's releases</a>.</em></p> <blockquote> <h2>v1.14.0</h2> <!-- raw HTML omitted --> <h2>✨ What's Changed</h2> <p>The main change in this release is that <code>verbose</code> and <code>print-hash</code> inputs are now on by default. This was contributed by <a href="https://github.com/whitequark"><code>@whitequark</code></a><a href="https://github.com/sponsors/whitequark">💰</a> in <a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/397">#397</a>.</p> <h2>📝 Docs</h2> <p><a href="https://github.com/woodruffw"><code>@woodruffw</code></a><a href="https://github.com/sponsors/woodruffw">💰</a> updated the mentions of PEP 740 to stop implying that it might be experimental (it hasn't been for quite a while!) in <a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/388">#388</a> and <a href="https://github.com/him2him2"><code>@him2him2</code></a><a href="https://github.com/sponsors/him2him2">💰</a> brushed up some grammar in the README and SECURITY docs via <a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/395">#395</a>.</p> <h2>🛠️ Internal Updates</h2> <p><a href="https://github.com/woodruffw"><code>@woodruffw</code></a><a href="https://github.com/sponsors/woodruffw">💰</a> bumped <code>sigstore</code> and <code>pypi-attestations</code> in the lock file (<a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/391">#391</a>) and <a href="https://github.com/webknjaz"><code>@webknjaz</code></a><a href="https://github.com/sponsors/webknjaz">💰</a> added infra for using type annotations in the project (<a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/381">#381</a>).</p> <h2>💪 New Contributors</h2> <ul> <li><a href="https://github.com/him2him2"><code>@him2him2</code></a> made their first contribution in <a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/395">#395</a></li> <li><a href="https://github.com/whitequark"><code>@whitequark</code></a> made their first contribution in <a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/397">#397</a></li> </ul> <p><strong>🪞 Full Diff</strong>: <a href="https://github.com/pypa/gh-action-pypi-publish/compare/v1.13.0...v1.14.0">https://github.com/pypa/gh-action-pypi-publish/compare/v1.13.0...v1.14.0</a></p> <p><strong>🧔♂️ Release Manager:</strong> <a href="https://github.com/sponsors/webknjaz"><code>@webknjaz</code></a> <a href="https://stand-with-ukraine.pp.ua">🇺🇦</a></p> <p><strong>🙏 Special Thanks</strong> to <a href="https://github.com/facutuesca"><code>@facutuesca</code></a><a href="https://github.com/sponsors/facutuesca">💰</a> and <a href="https://github.com/woodruffw"><code>@woodruffw</code></a><a href="https://github.com/sponsors/woodruffw">💰</a> for helping maintain this project when <a href="https://github.com/sponsors/webknjaz">I</a> can't!</p> <p><strong>💬 Discuss</strong> <a href="https://bsky.app/profile/webknjaz.me/post/3mivwsz3qzk2e">on Bluesky 🦋</a>, <a href="https://mastodon.social/@webknjaz/116363779997051422">on Mastodon 🐘</a> and <a href="https://github.com/pypa/gh-action-pypi-publish/discussions/404">on GitHub</a>.</p> <p><a href="https://github.com/sponsors/webknjaz"><img src="https://img.shields.io/badge/%40webknjaz-transparent?logo=githubsponsors&logoColor=%23EA4AAA&label=Sponsor&color=2a313c" alt="GH Sponsors badge" /></a></p> </blockquote> </details> <details> <summary>Commits</summary> <ul> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/cef221092ed1bacb1cc03d23a2d87d1d172e277b"><code>cef2210</code></a> Merge pull request <a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/397">#397</a> from whitequark/patch-1</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/b4595e2555a031e2fd6f0bbded4e7918eaa2724e"><code>b4595e2</code></a> Enable <code>verbose</code> and <code>print-hash</code> by default.</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/e2bab26859796ee5c3bf97b8f394ce1e6570e906"><code>e2bab26</code></a> Merge pull request <a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/395">#395</a> from him2him2/docs/fix-typos-and-grammar</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/7495c384ec7a0240a28e568e7ffc60af1629585d"><code>7495c38</code></a> docs: fix typos and grammar in README and SECURITY</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/03f86fee9ac21f854951f5c6e2a02c2a1324aec7"><code>03f86fe</code></a> Merge pull request <a href="https://redirect.github.com/pypa/gh-action-pypi-publish/issues/388">#388</a> from woodruffw-forks/ww/rm-experimental</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/4c78f1c53c55c528d8abd83df933ae92bd4c1d8c"><code>4c78f1c</code></a> Merge branch 'unstable/v1' into ww/rm-experimental</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/b5a6e8ba2611ad0c810f383eed9e6629eb0b3b2f"><code>b5a6e8b</code></a> deps: bump sigstore and pypi-attestations</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/a48a03e758da35722b0d159dae23e0440d0fcce2"><code>a48a03e</code></a> remove another experimental mention</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/8087a88a46924f78608905d7841a170e749524ce"><code>8087a88</code></a> action: remove a lingering mention of PEP 740 being experimental</li> <li><a href="https://github.com/pypa/gh-action-pypi-publish/commit/3317ede93a4981d0fc490510c6fcf8bf0e92ed05"><code>3317ede</code></a> 🧪 Integrate actionlint via pre-commit framework</li> <li>Additional commits viewable in <a href="https://github.com/pypa/gh-action-pypi-publish/compare/v1.13.0...v1.14.0">compare view</a></li> </ul> </details> <br /> [](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) --- <details> <summary>Dependabot commands and options</summary> <br /> You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show <dependency name> ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself) </details> Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
No description provided.