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

1from .activation_functions import ( 

2 SUPPORTED_ACTIVATIONS, 

3 ActivationFunction, 

4 gelu_fast, 

5 gelu_new, 

6 solu, 

7) 

8from .attribute_utils import get_nested_attr, set_nested_attr 

9from .components_utils import get_act_name 

10from .defaults_utils import ( 

11 USE_DEFAULT_VALUE, 

12 LocallyOverridenDefaults, 

13 override_or_use_default_value, 

14) 

15from .devices import ( 

16 ModelWithCfg, 

17 _MPS_MIN_SAFE_TORCH_VERSION, 

18 _mps_warned, 

19 _torch_version_tuple, 

20 get_device, 

21 move_to_and_update_config, 

22 warn_if_mps, 

23) 

24from .exploratory_utils import test_prompt 

25from .gpu_utils import print_gpu_mem 

26from .hf_utils import ( 

27 clear_huggingface_cache, 

28 download_file_from_hf, 

29 get_dataset, 

30 get_rotary_pct_from_config, 

31 keep_single_column, 

32 select_compatible_kwargs, 

33) 

34from .initialization_utils import ( 

35 NonlinearityType, 

36 calc_fan_in_and_fan_out, 

37 init_kaiming_normal_, 

38 init_kaiming_uniform_, 

39 init_xavier_normal_, 

40 init_xavier_uniform_, 

41) 

42from .library_utils import is_library_available 

43from .lm_utils import lm_accuracy, lm_cross_entropy_loss 

44from .logits_utils import sample_logits 

45from .matrix import ( 

46 composition_scores, 

47 get_matrix_corner, 

48) 

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 .slice import Slice, SliceInput 

63from .tensors import ( 

64 check_structure, 

65 filter_dict_by_prefix, 

66 get_corner, 

67 get_cumsum_along_dim, 

68 get_offset_position_ids, 

69 is_lower_triangular, 

70 is_square, 

71 remove_batch_dim, 

72 repeat_along_head_dimension, 

73 to_numpy, 

74 transpose, 

75) 

76from .tokenize_utils import ( 

77 get_attention_mask, 

78 get_input_with_manually_prepended_bos, 

79 get_tokenizer_with_bos, 

80 get_tokens_with_bos_removed, 

81 tokenize_and_concatenate, 

82)