Respect affine=False in SplitBatchNorm

[AHMETHAKANBEZIR1]

Description

SplitBatchNorm(..., affine=False) is accepted by its inherited BatchNorm constructor but fails during training with AttributeError: 'NoneType' object has no attribute 'repeat'. Both affine parameters are intentionally absent in this configuration.

Pass None to functional BatchNorm when the corresponding affine parameter is absent; repeat the parameters as before when present. Split assignment, running-stat updates, evaluation behavior and the default affine mode are unchanged. This addresses optional affine parameters, independently of #2080's input-layout fix.

Tests

• Added and ran regression tests.
• 24 cases across float32/float64, 1/2/4 splits, train/eval and affine on/off: clean master has 6 failures and 18 passing controls; the fix passes all 24. The final clean-base run asserted the imported module path in the same interpreter that ran pytest.
• Models suite: 279 passed, 102 skipped, 5 existing weight-norm deprecation warnings.
• Independent functional BatchNorm calls on the interleaved split samples verify output values, input/available affine gradients and running statistics. Evaluation uses the explicit mean of the split statistics.
• Locked Python 3.12.14 / torch 2.8.0 environment: lock check, global Ruff lint/format, full mypy (238 files), changed-file and commit hooks passed.

No GPU, distributed or older-version matrix was run locally. Constructor dtype/device behavior, untracked running-stat buffers, and cumulative momentum are outside this fix. No numerical algorithm or performance improvement is claimed.

Documentation

No public API or docstring changes. The unchanged all-files codespell false positive on FroSSL's valid ord="fro" is recorded in #2080; no rule was disabled.

AI disclosure: code, tests and description prepared autonomously with OpenAI Codex; no independent human review was performed.

SplitBatchNorm now passes None to functional BatchNorm when affine parameters are absent. It repeats the weight and bias across splits when present.

Added tests compare outputs, running statistics, and gradients across float32/float64, 1/2/4 splits, training/evaluation, and affine on/off.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 21f8db71-2a74-4197-93b8-fed67fa0c655

📥 Commits

Reviewing files that changed from the base of the PR and between 3cf8392 and 8e16c46.

📒 Files selected for processing: 2

• lightly/models/batchnorm.py
• tests/models/test_batchnorm_affine.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 5 remain after this review.

📝 Walkthrough

SplitBatchNorm now supports training when affine parameters are disabled. Tests compare its outputs, running statistics, and gradients with functional BatchNorm across several configurations.

Changes: SplitBatchNorm affine handling

Layer / File(s)	Summary
Handle missing affine weight lightly/models/batchnorm.py, tests/models/test_batchnorm_affine.py	The training path passes None when no weight exists and repeats the weight when it does. Parametrized tests compare outputs, running statistics, and gradients with functional BatchNorm.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix · Severity of issue fixed: Low

Merge Risk: ⚪ Minimal · up to 8e16c

The affine-disabled training fix has no identified issue that should delay merging after normal checks.

Pre-merge checks: 5

Passed checks: 5 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: ensuring that SplitBatchNorm respects affine=False.
Docstring Coverage		Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 2 files.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}