Skip to content

refactor(autogram): Create block-diagonal matrix using einsum#456

Merged
ValerianRey merged 1 commit intomainfrom
block-diagonal-is-einsum
Oct 15, 2025
Merged

refactor(autogram): Create block-diagonal matrix using einsum#456
ValerianRey merged 1 commit intomainfrom
block-diagonal-is-einsum

Conversation

@PierreQuinton
Copy link
Copy Markdown
Contributor

This removes a for loop and is therefore probably faster (especially on GPU).

@PierreQuinton PierreQuinton added cc: refactor Conventional commit type for any refactoring, not user-facing, and not typing or perf improvements package: autogram labels Oct 15, 2025
@codecov
Copy link
Copy Markdown

codecov bot commented Oct 15, 2025

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines Coverage Δ
src/torchjd/autogram/_engine.py 100.00% <100.00%> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@ValerianRey
Copy link
Copy Markdown
Contributor

It could be luck, but this seems to lead to a 1-5% speed improvement on GPU when batch_dim=None. It's much better IMO.

@ValerianRey ValerianRey changed the title refactor(autogram): Create block-diagonal matrix using einsum. refactor(autogram): Create block-diagonal matrix using einsum Oct 15, 2025
@ValerianRey ValerianRey merged commit 628e0b0 into main Oct 15, 2025
17 checks passed
@ValerianRey ValerianRey deleted the block-diagonal-is-einsum branch October 15, 2025 13:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cc: refactor Conventional commit type for any refactoring, not user-facing, and not typing or perf improvements package: autogram

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants