transformer_lens.utilities.parameter_swap module

Restore-guaranteed in-place parameter swapping.

torch.func.functional_call (and torch.nn.utils.stateless generally) fails to restore parameters on module trees that register the same submodule under more than one name: the tied-weight machinery swaps the single underlying slot once per alias, so the second swap stashes the override as the “original” and restoration installs the override permanently. TransformerBridge trees have exactly that shape — every replaced component is registered both in the wrapped HF tree and as a bridge submodule — so stateless reparametrization through a bridge silently corrupts it. This module provides the supported alternative: an in-place value swap whose restore is guaranteed by construction.

transformer_lens.utilities.parameter_swap.temporarily_swap_parameter(parameter: Any, new_value: Any) → Iterator[Parameter]

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

The parameter object is never replaced, so aliased registrations, optimizer references, hooks, and requires_grad state all stay intact. .grad is untouched. The restore runs in a finally block.

Parameters:

  • parameter – The live torch.nn.Parameter to modify.

  • new_value – Replacement values with the same shape and dtype. It may live on a different device; values are copied in.

Yields:

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

Raises:

  • TypeError – If parameter is not a torch.nn.Parameter or new_value is not a torch.Tensor.

  • ValueError – If shapes or dtypes differ — silent casts would make the swapped forward incomparable to the caller’s intent.