Coverage for transformer_lens/utilities/parameter_swap.py: 100%

1"""Restore-guaranteed in-place parameter swapping. 

2 

3``torch.func.functional_call`` (and ``torch.nn.utils.stateless`` generally) 

4fails to restore parameters on module trees that register the same submodule 

5under more than one name: the tied-weight machinery swaps the single underlying 

6slot once per alias, so the second swap stashes the override as the "original" 

7and restoration installs the override permanently. ``TransformerBridge`` trees 

8have exactly that shape — every replaced component is registered both in the 

9wrapped HF tree and as a bridge submodule — so stateless reparametrization 

10through a bridge silently corrupts it. This module provides the supported 

11alternative: an in-place value swap whose restore is guaranteed by construction. 

12""" 

13 

14from collections.abc import Iterator 

15from contextlib import contextmanager 

16from typing import Any 

17 

18import torch 

19 

20 

21@contextmanager 

22def temporarily_swap_parameter(parameter: Any, new_value: Any) -> Iterator[torch.nn.Parameter]: 

23 """Swap a parameter's value in place and restore it on exit, even on error. 

24 

25 The parameter object is never replaced, so aliased registrations, optimizer 

26 references, hooks, and ``requires_grad`` state all stay intact. ``.grad`` is 

27 untouched. The restore runs in a ``finally`` block. 

28 

29 Args: 

30 parameter: The live ``torch.nn.Parameter`` to modify. 

31 new_value: Replacement values with the same shape and dtype. It may live 

32 on a different device; values are copied in. 

33 

34 Yields: 

35 The same parameter, holding ``new_value`` for the duration of the block. 

36 

37 Raises: 

38 TypeError: If ``parameter`` is not a ``torch.nn.Parameter`` or 

39 ``new_value`` is not a ``torch.Tensor``. 

40 ValueError: If shapes or dtypes differ — silent casts would make the 

41 swapped forward incomparable to the caller's intent. 

42 """ 

43 # Any-typed params keep these manual checks live under the project's 

44 # beartype instrumentation; the docstring states the real types. 

45 if not isinstance(parameter, torch.nn.Parameter): 

46 raise TypeError(f"parameter must be a torch.nn.Parameter; got {type(parameter).__name__}") 

47 if isinstance(new_value, torch.nn.Parameter) or not isinstance(new_value, torch.Tensor): 

48 raise TypeError(f"new_value must be a plain torch.Tensor; got {type(new_value).__name__}") 

49 if new_value.shape != parameter.shape: 

50 raise ValueError( 

51 f"new_value shape {tuple(new_value.shape)} must match parameter shape " 

52 f"{tuple(parameter.shape)}" 

53 ) 

54 if new_value.dtype != parameter.dtype: 

55 raise ValueError( 

56 f"new_value dtype {new_value.dtype} must match parameter dtype {parameter.dtype}" 

57 ) 

58 original = parameter.detach().clone() 

59 try: 

60 with torch.no_grad(): 

61 parameter.copy_(new_value) 

62 yield parameter 

63 finally: 

64 with torch.no_grad(): 

65 parameter.copy_(original)