Skip to content

fix: remove unconditional inplace_mask_diag that overrides owner self…#2597

Open
greatjourney589 wants to merge 2 commits intoopentensor:devnet-readyfrom
greatjourney589:fix/epoch-owner-diagonal-mask
Open

fix: remove unconditional inplace_mask_diag that overrides owner self…#2597
greatjourney589 wants to merge 2 commits intoopentensor:devnet-readyfrom
greatjourney589:fix/epoch-owner-diagonal-mask

Conversation

@greatjourney589
Copy link
Copy Markdown

@greatjourney589 greatjourney589 commented Apr 17, 2026

Description

Remove a stray unconditional inplace_mask_diag(&mut weights) call in epoch_dense_mechanism that immediately followed the owner self-weight exception block. When an owner_uid is present, the if branch correctly calls inplace_mask_diag_except_index to preserve the owner's self-weight on the diagonal. However, the redundant call afterward would unconditionally zero the entire diagonal — silently overriding and negating that exception on every epoch.

Related Issue(s)

  • Closes #[issue number]

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Documentation update
  • Other (please describe):

Breaking Change

N/A — this fix restores intended behavior and does not alter any public interfaces.

Checklist

  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas
  • I have run ./scripts/fix_rust.sh to ensure my code is formatted and linted correctly
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • Any dependent changes have been merged and published in downstream modules

Screenshots (if applicable)

Please include any relevant screenshots or GIFs that demonstrate the changes made.

Additional Notes

The removed line appears to be a leftover from before the owner self-weight exception was introduced. The if let Some(owner_uid) block already handles both branches (inplace_mask_diag_except_index for the owner case, inplace_mask_diag for the general case), making the unconditional call after the block both redundant and harmful.

@greatjourney589
Copy link
Copy Markdown
Author

Hi @open-junius @sam0x17 @gztensor @JohnReedV @camfairchild
I discovered a potential issue during codebase exploration. I’d appreciate your review of this PR.

@camfairchild
Copy link
Copy Markdown
Contributor

camfairchild commented Apr 20, 2026

utACK

Thanks for the contribution

Copy link
Copy Markdown
Contributor

@camfairchild camfairchild left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please write a test for this also

@greatjourney589 greatjourney589 force-pushed the fix/epoch-owner-diagonal-mask branch from b9ddcb3 to 4a3d782 Compare April 20, 2026 11:54
@greatjourney589
Copy link
Copy Markdown
Author

Please write a test for this also

A corresponding test has been added. Thank you for catching that.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants