Fix BatchLM vmap+jacfwd failures caused by in-place tensor mutations

[Copilot]

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.

# 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

[codecov]

Codecov Report

❌ Patch coverage is 92.85714% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.54%. Comparing base (5691ab1) to head (a6dc32f).
⚠️ Report is 1 commits behind head on dev.

Files with missing lines	Patch %	Lines
astrophot/backend_obj.py	87.50%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##              dev     #301      +/-   ##
==========================================
+ Coverage   91.23%   91.54%   +0.30%     
==========================================
  Files         113      113              
  Lines        6198     6194       -4     
==========================================
+ Hits         5655     5670      +15     
+ Misses        543      524      -19     

Flag	Coverage Δ
unittests	91.54% <92.85%> (+0.30%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Fixes BatchLM failures under vmap(jacfwd(...)) by removing in-place tensor mutations (and some bool-mask dynamic-shape patterns) from the Torch execution path, and adds regression tests to cover the affected integration/sampling configurations.

Changes:

• Updated Torch backend indexed write/add helpers to use functional index_put for integer tensor indices (forward-mode AD compatible).
• Reworked bool-mask accumulation paths in brightness mixins and spline extrapolation to use where-based masking (static shapes under vmap).
• Added BatchLM test coverage across integrate_mode, sampling_mode, rotated WCS/CD, and Poisson likelihood.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
astrophot/backend_obj.py	Switches Torch indexed updates to index_put for LongTensor indices; removes unused and_at_indices.
astrophot/models/mixins/brightness.py	Replaces bool-mask indexed updates with where masking to keep shapes static under vmap+jacfwd.
astrophot/models/func/spline.py	Replaces extrapolation-time masked fills with where to avoid bool-index mutation/shape issues.
tests/test_fit.py	Adds targeted regression tests ensuring BatchLM runs across integration/sampling modes and likelihood variants.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.