[CeliaBenquet]

Unified solvers (cebra/solver/unified_session.py)

• UnifiedSolverxCEBRA (registered as "unified-xcebra"): unified CEBRA with the xCEBRA Jacobian regularizer. It uses a single objective; it builds on the multi-objective base class only for the regularizer and its weight schedule.
• solver.aligned_inputs(...) builds the unified model input for a session. One shared implementation replaces the copies we had in previous version.
• transform(), for both unified solvers:
  • The previous version padded every batch separately. Rows near batch edges therefore saw padding instead of their neighbours; this is fixed. Also no longer fails for sessions shorter than twice batch_size.
  • random_ties added, to make sure if there are nearest-neighbor ties, this is sampled uniformly (until now, first one was the default).
• split_outputs is removed from the unified transform(); it always returns the concatenated embedding, as multiple objectives was not implemented.

Feature attribution maps (cebra.attribution)

A new module, cebra/attribution/feature_map.py, computes how each input channel drives each behavioural or stimulus feature through the embedding, for models trained with the Jacobian regularizer. It exposes three functions:

• fit_feature_readout(embedding, features, alpha=1.0, train_fraction=None, train_index=None): fits a ridge readout W (D × F) from the embedding to the features.
• feature_attribution_map(solver, inputs, W, *, labels=None, session_id=None, indices=None, output_slice=None, average=True, random_ties=False, seed=0, ...): the signed map M = Jᵀ W, with J the Jacobian of the embedding with respect to the input window. average=True returns the time mean (N × F). It accepts single-session solvers, including one head of a subspace model via output_slice, and unified solvers. It accepts only solvers trained with the Jacobian regularizer.
• average_feature_maps(maps): the mean of maps from several models. Each model needs its own W.

Cleanup and fixes

• Dependencies and debug output:
  • Removed the hard wandb import and the stray print calls
  • torch_brain is imported only when the Perich loader needs it, and perich.py has its licence header.
• Checkpoint loading:
  • Solver.load defaults to the safe path (weights_only=None, which is True on torch ≥ 2.6).
  • The allowed types now include numpy scalars and dtypes, so checkpoints saved under numpy 1 or 2 load safely.
• Removed unused code: expand_index_forward and ForwardTensorDataset.
• Masking: simplified by removing the combined-mask and noise-filling options.
• Docs: API pages added for the unified solvers, masking, the unified datasets and loaders, and feature_map.

[Copilot]

🟡 Changes recommended

Unresolved critical defects can leak masks across datasets, make subsampled inference depend on batch size, and break learning-rate scheduler progression and checkpoint restoration.

8 open issues

• Copy inherited masks before appending list-format configuration · New
• Build NumPy dtype allowlist compatibly across versions · New
• Save and restore scheduler state in checkpoints · New
• Advance schedulers in every solver update path · New
• Preserve whole-session inference for resampling models · New
• Avoid recomputation for training-mode models · New
• Keep unified sampling out of ContrastiveMultiObjectiveLoader · New
• Use configured time_delta for unified sampling · New

What changed in this PR

This PR extends CEBRA’s PyTorch APIs with regularized unified-session training and feature attribution maps, while improving unified inference and checkpoint handling.

Changes:

• Add unified xCEBRA training and signed feature maps linking input channels to features through embeddings.
• Rework unified inference and nearest-label sampling, including reproducible tie handling and batching improvements.
• Add scheduler support, masking and dataset utilities, regression tests, and API documentation.

File	Description
tests/​test_solver.py	Test unified inference, alignment, and checkpoint loading.
tests/​test_models.py	Test Jacobian regularizer values and gradients.
tests/​test_integration_xcebra.py	Move embeddings to CPU for NumPy comparisons.
tests/​test_feature_map.py	Test readouts, feature maps, and solver dispatch.
tests/​test_distributions.py	Test tie sampling and optimized nearest-label search.
tests/​test_data_masking.py	Test session-block masking.
tests/​test_attribution.py	Test optimized attribution and output batching.
docs/​source/​api/​pytorch/​solvers.rst	Document unified solvers and link multi-objective APIs.
docs/​source/​api/​pytorch/​data.rst	Add unified-data and masking references.
docs/​source/​api/​pytorch/​attribution.rst	Expose feature-map API documentation.
cebra/​solver/​unified_session.py	Share alignment logic, revise inference, and add unified xCEBRA.
cebra/​solver/​multiobjective.py	Link unified solvers and normalize logged NumPy scalars.
cebra/​solver/​base.py	Add checkpoint allowlisting and learning-rate scheduler support.
cebra/​models/​multiobjective.py	Remove debug prints.
cebra/​models/​jacobian_regularizer.py	Create projection tensors on the input device.
cebra/​distributions/​multisession.py	Add device-aware sampling and optimized tie handling.
cebra/​distributions/​__init__.py	Export UnifiedSampler.
cebra/​datasets/​allen/​single_session_ca.py	Add configurable cortex and neuron-type loading.
cebra/​datasets/​__init__.py	Whitespace-only cleanup.
cebra/​data/​multiobjective.py	Add unified contrastive loading and sampling configuration.
cebra/​data/​multi_session.py	Move sampled indices to the loader device.
cebra/​data/​masking.py	Add list-based configuration and session-block masks.
cebra/​data/​datasets.py	Extend device transfers and add unified xCEBRA batches.
cebra/​attribution/​feature_map.py	Implement ridge readouts and signed feature attribution maps.
cebra/​attribution/​attribution_models.py	Add optimized Jacobian batching and lazy optional imports.
cebra/​attribution/​_jacobian.py	Add output-dimension chunking for Jacobian computation.
cebra/​attribution/​__init__.py	Export the feature-map functions.

_{Review effort: Balanced}