[emerardd]

Description

Follow-up to #1526: the native-autograd normalization path handles replacement tensors, but mistakes same-object in-place edits for observation hooks. For example, returning value.mul_(2) from hook_normalized or hook_scale has no effect, while returning value * 2 changes the result and emits the expected edit-fallback warning.

This change detects tensor writes as well as replacements at both intermediate hooks. A scale edit recomputes the normalized activation before hook_normalized; either edit selects reconstruction from the hooked values. Mutations also propagate when the callback returns None, consistent with the existing Python normalization path.

Hooks receive grad-connected clones so an in-place scale edit does not invalidate the original square-root output saved by autograd. Normal execution uses version counters for mutation detection. Inference tensors have no version counter, so inference-mode callbacks use an exact, NaN-aware value snapshot instead. Pure observers still select HF's native output, preserving exact forward and backward numerics, and no-hook execution avoids cloning.

The existing edit-fallback and backward-hook warnings remain unchanged. No public signatures or model-loading behavior change. The native-path docstring now includes in-place edits.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Validation

Before the fix, the added edit/observation matrix had 36 failed / 6 passed. Coverage includes scale/normalized hooks, returned tensors and None, LayerNorm/RMSNorm/offset RMSNorm, grad/no-grad/inference execution, input and affine-parameter gradients, exact observation parity, and NaN/Inf observations under inference mode.

A real config-built tiny Llama also checks layer-internal and final normalization interventions, HF/Bridge and cache parity, equivalent replacement edits and hook cleanup. No pretrained weights are downloaded.

• Final five-file normalization/offset/Llama regression surface: 93 passed.
• Full mypy .: 389 source files passed, including the final incremental check.
• Pinned pycln/isort/Black checks on all three changed files and git diff --check: passed.

The complete unit suite and GPU/MPS hardware tests were not run.

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (native-path docstring)
• 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 (affected tests run; full suite not run)
• I have not rewritten tests relating to key interfaces which would affect backward compatibility