transformer_lens.conversion_utils package

Subpackages

• transformer_lens.conversion_utils.conversion_steps package
  • Submodules
    • transformer_lens.conversion_utils.conversion_steps.arithmetic_tensor_conversion module
      • ArithmeticTensorConversion
      • OperationTypes
    • transformer_lens.conversion_utils.conversion_steps.attention_auto_conversion module
      • AttentionAutoConversion
    • transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion module
      • BaseTensorConversion
    • transformer_lens.conversion_utils.conversion_steps.callable_tensor_conversion module
      • CallableTensorConversion
    • transformer_lens.conversion_utils.conversion_steps.chain_tensor_conversion module
      • ChainTensorConversion
    • transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion module
      • RearrangeTensorConversion
    • transformer_lens.conversion_utils.conversion_steps.repeat_tensor_conversion module
      • RepeatTensorConversion
    • transformer_lens.conversion_utils.conversion_steps.split_tensor_conversion module
      • SplitTensorConversion
    • transformer_lens.conversion_utils.conversion_steps.tensor_conversion_set module
      • TensorConversionSet
    • transformer_lens.conversion_utils.conversion_steps.ternary_tensor_conversion module
      • TernaryTensorConversion
    • transformer_lens.conversion_utils.conversion_steps.transpose_tensor_conversion module
      • TransposeTensorConversion
    • transformer_lens.conversion_utils.conversion_steps.zeros_like_conversion module
      • ZerosLikeConversion
  • Module contents
    • ArithmeticTensorConversion
      • ArithmeticTensorConversion.handle_conversion()
      • ArithmeticTensorConversion.revert()
    • AttentionAutoConversion
      • AttentionAutoConversion.__init__()
      • AttentionAutoConversion.clear_state()
      • AttentionAutoConversion.get_conversion_info()
      • AttentionAutoConversion.handle_conversion()
      • AttentionAutoConversion.revert_conversion()
    • BaseTensorConversion
      • BaseTensorConversion.convert()
      • BaseTensorConversion.handle_conversion()
      • BaseTensorConversion.revert()
    • CallableTensorConversion
      • CallableTensorConversion.handle_conversion()
    • ChainTensorConversion
      • ChainTensorConversion.__init__()
      • ChainTensorConversion.handle_conversion()
      • ChainTensorConversion.revert()
    • RearrangeTensorConversion
      • RearrangeTensorConversion.handle_conversion()
      • RearrangeTensorConversion.revert()
    • RepeatTensorConversion
      • RepeatTensorConversion.handle_conversion()
    • SplitTensorConversion
      • SplitTensorConversion.__init__()
      • SplitTensorConversion.handle_conversion()
    • TensorConversionSet
      • TensorConversionSet.get_component()
      • TensorConversionSet.get_conversion_action()
      • TensorConversionSet.handle_conversion()
      • TensorConversionSet.process_conversion()
      • TensorConversionSet.process_conversion_action()
    • TernaryTensorConversion
      • TernaryTensorConversion.find_context_field()
      • TernaryTensorConversion.handle_conversion()
      • TernaryTensorConversion.handle_fallback_conversion()
      • TernaryTensorConversion.handle_primary_conversion()
    • TransposeTensorConversion
      • TransposeTensorConversion.handle_conversion()
      • TransposeTensorConversion.revert()
    • ZerosLikeConversion
      • ZerosLikeConversion.handle_conversion()
• transformer_lens.conversion_utils.helpers package
  • Submodules
    • transformer_lens.conversion_utils.helpers.find_property module
      • find_property()
    • transformer_lens.conversion_utils.helpers.merge_quantiziation_fields module
      • merge_quantization_fields()
  • Module contents
    • find_property()
    • merge_quantization_fields()

Submodules

• transformer_lens.conversion_utils.hook_conversion_utils module
  • get_weight_conversion_field_set()
  • model_info_cfg()
• transformer_lens.conversion_utils.param_processing_conversion module
  • ParamProcessingConversion
    • ParamProcessingConversion.convert()
    • ParamProcessingConversion.revert()

Module contents

Model bridge conversion utilities.

This module contains utilities for converting between different model architectures.

class transformer_lens.conversion_utils.TensorConversionSet(fields: dict[str, Any])

Bases: BaseTensorConversion

get_component(model: Any, name: str) → Any

Get a component from the model using the field mapping.

Parameters:

• model – The model to get the component from.

• name – The name of the component to get.

Returns:

The requested component.

get_conversion_action(field: str) → BaseTensorConversion

handle_conversion(input_value: Any, *full_context: Any) → dict[str, Any]

process_conversion(input_value: Any, remote_field: str, conversion: BaseTensorConversion, *full_context: Any) → Any

process_conversion_action(input_value: Any, conversion_details: Any, *full_context: Any) → Any

transformer_lens.conversion_utils.convert_tl_checkpoint(state_dict: dict[str, Tensor], cfg: TransformerBridgeConfig) → dict[str, Tensor]

Convert a legacy TL-property-format state dict to the key/tensor format TransformerBridge.boot_native(cfg).load_state_dict accepts.

Parameters:

• state_dict – A state dict in the old HookedTransformer convention (e.g. from HookedTransformer.state_dict()), with keys like "blocks.0.attn.W_Q" and per-head tensor shapes.

• cfg – The config the checkpoint was trained/saved under. Used both to reshape per-head attention weights and to validate that the checkpoint’s per-head shapes actually match this cfg — a mismatched cfg would otherwise silently mis-group heads without ever tripping a shape error, since d_model == n_heads * d_head holds for any wrong factoring too.

Returns:

A state dict with modern bridge keys (e.g. "blocks.0.attn.q.weight") and flat nn.Linear-oriented tensor shapes, ready for bridge.load_state_dict(converted, strict=True).