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

1from .activation_functions import ( 

2 SOFTCAP_DISABLED, 

3 SUPPORTED_ACTIVATIONS, 

4 ActivationFunction, 

5 apply_softcap, 

6 gelu_fast, 

7 gelu_new, 

8 softcap_enabled, 

9 solu, 

10) 

11from .attribute_utils import get_nested_attr, set_nested_attr 

12from .components_utils import get_act_name 

13from .defaults_utils import ( 

14 USE_DEFAULT_VALUE, 

15 LocallyOverridenDefaults, 

16 override_or_use_default_value, 

17) 

18from .devices import ( 

19 ModelWithCfg, 

20 _MPS_MIN_SAFE_TORCH_VERSION, 

21 _mps_warned, 

22 _torch_version_tuple, 

23 get_device, 

24 move_to_and_update_config, 

25 warn_if_mps, 

26) 

27from .exploratory_utils import test_prompt 

28from .gpu_utils import print_gpu_mem 

29from .hf_utils import ( 

30 clear_huggingface_cache, 

31 download_file_from_hf, 

32 get_dataset, 

33 get_rotary_pct_from_config, 

34 keep_single_column, 

35 select_compatible_kwargs, 

36) 

37from .initialization_utils import ( 

38 NonlinearityType, 

39 calc_fan_in_and_fan_out, 

40 init_kaiming_normal_, 

41 init_kaiming_uniform_, 

42 init_xavier_normal_, 

43 init_xavier_uniform_, 

44) 

45from .library_utils import is_library_available 

46from .lm_utils import lm_accuracy, lm_cross_entropy_loss 

47from .logits_utils import logits_to_df, sample_logits 

48from .matrix import ( 

49 composition_scores, 

50 get_matrix_corner, 

51) 

52 

53# Re-export multi-GPU helpers here (devices.py must not import multi_gpu directly) 

54from .multi_gpu import ( 

55 calculate_available_device_cuda_memory, 

56 count_unique_devices, 

57 determine_available_memory_for_available_devices, 

58 find_embedding_device, 

59 get_best_available_cuda_device, 

60 get_best_available_device, 

61 get_device_for_block_index, 

62 resolve_device_map, 

63 sort_devices_based_on_available_memory, 

64) 

65from .slice import Slice, SliceInput 

66from .tensors import ( 

67 check_structure, 

68 filter_dict_by_prefix, 

69 get_corner, 

70 get_cumsum_along_dim, 

71 get_offset_position_ids, 

72 is_lower_triangular, 

73 is_square, 

74 remove_batch_dim, 

75 repeat_along_head_dimension, 

76 to_numpy, 

77 transpose, 

78) 

79from .tokenize_utils import ( 

80 get_attention_mask, 

81 get_input_with_manually_prepended_bos, 

82 get_tokenizer_with_bos, 

83 get_tokens_with_bos_removed, 

84 tokenize_and_concatenate, 

85) 

86from .typed_module_list import TypedModuleList