Fix JAX portability and extreme-scale numerical regressions

[FlorianPfaff]

Summary

Fix five concrete production defects exposed by the existing cross-backend regression suite, primarily on JAX:

• treat nested boolean lists as boolean masks in the JAX assignment compatibility layer instead of misclassifying them as coordinate lists;
• normalize relaxed-S3F grid masses with split scale division so valid ratios near the backend floating-point maximum are preserved;
• apply the same stable normalization to SO3ProductParticleFilter weights;
• validate and trace-normalize covariance-shape features after positive rescaling, avoiding raw-scale eigensolver/reciprocal failures for extreme finite PSD covariances;
• remove in-place masked assignments from IMM log-weight construction so zero likelihood/prior mass is handled on immutable JAX arrays.

Evidence

The repository's existing Python 3.13 JAX test artifact reproduces the affected regressions on current main:

• tests/backend_support/test_jax_assignment_contract.py::test_jax_assignment_accepts_list_and_boolean_indices
• tests/filters/test_relaxed_s3f_grid_mass_overflow.py::{test_large_finite_grid_values_normalize_without_overflow,test_circular_mean_is_invariant_to_large_common_weight_scale}
• tests/filters/test_so3_product_particle_filter_extreme_weights.py::{test_constructor_normalizes_extreme_finite_weights,test_set_particles_normalizes_extreme_finite_weights}
• tests/test_pairwise_covariance_symmetrization_overflow.py::test_shape_components_preserve_extreme_finite_covariances
• tests/filters/test_imm_positive_infinite_log_likelihood.py::test_negative_infinite_log_likelihood_remains_zero_mass

No new tests are needed because these failures are already encoded as focused regressions.

Numerical rationale

For the extreme-weight paths, replacing x / scale with (x / sqrt(scale)) / sqrt(scale) is algebraically identical but avoids JAX lowering the former through a reciprocal that can underflow near the largest representable floating-point value. PyRecEst already uses this stabilization in the IMM's nonnegative-weight normalizer.

For covariance PSD/shape operations, positive scalar rescaling preserves inertia and trace-normalized shape while keeping the eigensolver and divisions in a well-conditioned numerical range.

Scope

• based directly on current main at baac3b1736bb8e81eafc1e3f69e10104636103b9;
• 5 commits ahead, 0 behind at PR creation;
• 5 production files changed;
• no overlap found with currently open bug-fix PRs Fix range-bearing Jacobian per-call overrides #5354–Align covariance logdet regression with PSD validation #5370.

[github-actions]

✅MegaLinter analysis: Success

Descriptor	Linter	Files	Fixed	Errors	Max errors	Warnings	Elapsed time
✅ COPYPASTE	jscpd	yes		no		no	8.02s
✅ JSON	prettier	7	0	0		0	0.75s
✅ JSON	v8r	7		0		0	3.83s
✅ MARKDOWN	markdownlint	68	0	0		0	2.03s
✅ MARKDOWN	markdown-table-formatter	68	0	0		0	0.42s
✅ PYTHON	black	1988	16	0		0	58.3s
✅ PYTHON	isort	1988	28	0		0	2.8s
✅ REPOSITORY	betterleaks	yes		no		no	1.52s
✅ REPOSITORY	checkov	yes		no		no	34.06s
✅ REPOSITORY	git_diff	yes		no		no	0.26s
✅ REPOSITORY	secretlint	yes		no		no	99.05s
✅ REPOSITORY	syft	yes		no		no	3.14s
✅ REPOSITORY	trivy-sbom	yes		no		no	2.52s
✅ YAML	prettier	11	0	0		0	0.85s
✅ YAML	v8r	11		0		0	13.31s
✅ YAML	yamllint	11		0		0	0.79s

Notices

⚠️ Your configuration references items that have been removed from MegaLinter and are ignored: MAKEFILE_CHECKMAKE. See Removed linters to find their replacements.

See detailed reports in MegaLinter artifacts

Your project could benefit from a custom flavor, which would allow you to run only the linters you need, and thus improve runtime performances. (Skip this info by defining FLAVOR_SUGGESTIONS: false)

• Documentation: Custom Flavors
• Command: npx mega-linter-runner@10.0.0 --custom-flavor-setup --custom-flavor-linters PYTHON_BLACK,PYTHON_ISORT,COPYPASTE_JSCPD,JSON_V8R,JSON_PRETTIER,MARKDOWN_MARKDOWNLINT,MARKDOWN_MARKDOWN_TABLE_FORMATTER,REPOSITORY_CHECKOV,REPOSITORY_GIT_DIFF,REPOSITORY_BETTERLEAKS,REPOSITORY_SECRETLINT,REPOSITORY_SYFT,REPOSITORY_TRIVY_SBOM,YAML_PRETTIER,YAML_YAMLLINT,YAML_V8R

Show us your support by starring ⭐ the repository