Skip to content

fix(interop): five bugs found in conversion audit#213

Merged
chaoming0625 merged 1 commit into
mainfrom
worktree-interop-audit
Jun 11, 2026
Merged

fix(interop): five bugs found in conversion audit#213
chaoming0625 merged 1 commit into
mainfrom
worktree-interop-audit

Conversation

@chaoming0625

Copy link
Copy Markdown
Collaborator

Summary

  • Critical: nnx Conv import silently dropped input_dilation, producing numerically wrong results. Added the same guard as the linen adapter.
  • High: Norm layer conversions (LayerNorm, RMSNorm, GroupNorm) crashed with AttributeError when all affine parameters were disabled (use_scale=False, use_bias=False). Fixed by extracting feature count from framework metadata (m.num_features, m.shape, layer.in_size) instead of from affine parameter arrays.
  • High: bst_set_norm crashed when the norm state attribute was None (no affine). Added early return.
  • Medium: bst_set_batchnorm wrote None values into the weight dict. Now omits None-valued keys.
  • Medium: lookup_export rebuilt a full dict from _EXPORT on every call. Replaced with direct __mro__-walk over the existing dict.

Test plan

  • All 70 existing interop tests pass (no regressions)
  • 8 new edge-case tests added and passing:
    • nnx LayerNorm/GroupNorm with no affine params (import + export round-trip)
    • nnx LayerNorm scale-only and bias-only (import + export round-trip)
    • equinox LayerNorm/GroupNorm/RMSNorm with no affine params (import + export round-trip)
    • nnx Conv with input_dilation != 1 raises ConversionError

🤖 Generated with Claude Code

- Add missing input_dilation guard to nnx Conv import (silent data corruption)
- Fix norm num extraction to use framework metadata instead of affine params
  that may be None (crashes on LayerNorm/RMSNorm/GroupNorm without affine)
- Fix bst_set_norm to early-return when both scale and offset are None
- Fix bst_set_batchnorm to omit None-valued keys from weight dict
- Fix lookup_export to avoid O(N) dict rebuild on every call

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Sorry @chaoming0625, you have reached your weekly rate limit of 500000 diff characters.

Please try again later or upgrade to continue using Sourcery

@chaoming0625 chaoming0625 merged commit f726f07 into main Jun 11, 2026
6 checks passed
@codecov

codecov Bot commented Jun 11, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 78.68852% with 13 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
brainstate/interop/_frameworks/_linen.py 52.38% 4 Missing and 6 partials ⚠️
brainstate/interop/_common.py 75.00% 0 Missing and 2 partials ⚠️
brainstate/interop/_frameworks/_nnx.py 93.75% 0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

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.

1 participant