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 .tensors import ( 

65 check_structure, 

66 filter_dict_by_prefix, 

67 get_corner, 

68 get_cumsum_along_dim, 

69 get_offset_position_ids, 

70 is_lower_triangular, 

71 is_square, 

72 remove_batch_dim, 

73 repeat_along_head_dimension, 

74 to_numpy, 

75 transpose, 

76) 

77from .tokenize_utils import ( 

78 get_attention_mask, 

79 get_input_with_manually_prepended_bos, 

80 get_tokenizer_with_bos, 

81 get_tokens_with_bos_removed, 

82 tokenize_and_concatenate, 

83) 

84from .typed_module_list import TypedModuleList