Fix 1130 and 1102

tests/unit/test_utils.py

@@ -165,6 +165,18 @@ def test_override_or_use_default_value():
     assert utils.override_or_use_default_value(default_flag=False) == False
 
 
+def test_is_library_available():
+    # Checking libraries that are definitely included in dependencies
+    assert utils.is_library_available("torch") is True
+    assert utils.is_library_available("numpy") is True
+
+    # Check the standard library (sys is always loaded)
+    assert utils.is_library_available("sys") is True
+
+    # Checking a library that obviously doesn't exist
+    assert utils.is_library_available("completely_fake_library_name_123") is False
+
+
 class TestAttentionMask:
     prompts = [
         "Hello world!",

transformer_lens/HookedTransformer.py

@@ -2010,7 +2010,7 @@ def set_use_attn_result(self, use_attn_result: bool):
 
     def set_use_split_qkv_input(self, use_split_qkv_input: bool):
         """
-        Toggles whether to allow editing of inputs to each attention head.
+        Toggles whether to allow editing of the separate Q, K, and V inputs to each attention head.
         """
         self.cfg.use_split_qkv_input = use_split_qkv_input
 

transformer_lens/train.py

@@ -9,13 +9,13 @@
 
 import torch
 import torch.optim as optim
-import wandb
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, Dataset
 from tqdm.auto import tqdm
 
 from transformer_lens import utils
 from transformer_lens.HookedTransformer import HookedTransformer
+from transformer_lens.utils import is_library_available
 
 
 @dataclass
@@ -74,9 +74,16 @@ def train(
     Returns:
         The trained model
     """
+
     torch.manual_seed(config.seed)
     model.train()
+
     if config.wandb:
+        if not is_library_available("wandb"):
+            raise ImportError("Wandb is not available")
+
+        import wandb
+
         if config.wandb_project_name is None:
             config.wandb_project_name = "easy-transformer"
         wandb.init(project=config.wandb_project_name, config=vars(config))

transformer_lens/utils.py

@@ -6,11 +6,13 @@
 from __future__ import annotations
 
 import collections.abc
+import importlib.util
 import inspect
 import json
 import os
 import re
 import shutil
+import sys
 from copy import deepcopy
 from typing import Any, List, Optional, Tuple, Union, cast
 
@@ -34,6 +36,15 @@
 USE_DEFAULT_VALUE = None
 
 
+def is_library_available(name: str) -> bool:
+    """
+    Checks if a library is installed in the current environment without importing it.
+    Prevents crash or segmentation fault.
+    """
+
+    return name in sys.modules or importlib.util.find_spec(name) is not None
+
+
 def select_compatible_kwargs(
     kwargs_dict: dict[str, Any], callable: collections.abc.Callable
 ) -> dict[str, Any]: