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 _MPS_MIN_SAFE_TORCH_VERSION, 

20 ModelWithCfg, 

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 composition_scores, get_matrix_corner 

49 

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

51from .multi_gpu import ( 

52 calculate_available_device_cuda_memory, 

53 count_unique_devices, 

54 determine_available_memory_for_available_devices, 

55 find_embedding_device, 

56 get_best_available_cuda_device, 

57 get_best_available_device, 

58 get_device_for_block_index, 

59 resolve_device_map, 

60 sort_devices_based_on_available_memory, 

61) 

62from .parameter_swap import temporarily_swap_parameter 

63from .slice import Slice, SliceInput 

64from .statistics import ( 

65 bootstrap_ci, 

66 derive_generator, 

67 sign_flip_permutation_pvalue, 

68 standard_error, 

69) 

70from .tensors import ( 

71 check_structure, 

72 filter_dict_by_prefix, 

73 get_corner, 

74 get_cumsum_along_dim, 

75 get_offset_position_ids, 

76 is_lower_triangular, 

77 is_square, 

78 remove_batch_dim, 

79 repeat_along_head_dimension, 

80 to_numpy, 

81 transpose, 

82) 

83from .tokenize_utils import ( 

84 get_attention_mask, 

85 get_input_with_manually_prepended_bos, 

86 get_tokenizer_with_bos, 

87 get_tokens_with_bos_removed, 

88 tokenize_and_concatenate, 

89) 

90from .typed_module_list import TypedModuleList