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

1"""Benchmark utilities for TransformerBridge testing. 

2 

3This module provides reusable benchmark functions for validating TransformerBridge 

4against raw HuggingFace models and pre-computed reference data (golden fixtures). 

5""" 

6 

7from transformer_lens.benchmarks.activation_cache import ( 

8 benchmark_activation_cache, 

9 benchmark_run_with_cache, 

10) 

11from transformer_lens.benchmarks.backward_gradients import ( 

12 benchmark_backward_hooks, 

13 benchmark_critical_backward_hooks, 

14 benchmark_gradient_computation, 

15) 

16from transformer_lens.benchmarks.forward_pass import ( 

17 benchmark_forward_pass, 

18 benchmark_logits_equivalence, 

19 benchmark_loss_equivalence, 

20) 

21from transformer_lens.benchmarks.generation import ( 

22 benchmark_generation, 

23 benchmark_generation_with_kv_cache, 

24 benchmark_multiple_generation_calls, 

25) 

26from transformer_lens.benchmarks.hook_registration import ( 

27 benchmark_critical_forward_hooks, 

28 benchmark_forward_hooks, 

29 benchmark_gated_hooks_fire, 

30 benchmark_hook_functionality, 

31 benchmark_hook_registry, 

32) 

33from transformer_lens.benchmarks.main_benchmark import run_benchmark_suite 

34from transformer_lens.benchmarks.text_quality import benchmark_text_quality 

35from transformer_lens.benchmarks.utils import ( 

36 BenchmarkResult, 

37 BenchmarkSeverity, 

38 PhaseReferenceData, 

39) 

40from transformer_lens.benchmarks.weight_processing import benchmark_weight_modification 

41 

42__all__ = [ 

43 # Main benchmark runner 

44 "run_benchmark_suite", 

45 # Result types 

46 "BenchmarkResult", 

47 "BenchmarkSeverity", 

48 "PhaseReferenceData", 

49 # Forward pass benchmarks 

50 "benchmark_forward_pass", 

51 "benchmark_logits_equivalence", 

52 "benchmark_loss_equivalence", 

53 # Hook benchmarks 

54 "benchmark_forward_hooks", 

55 "benchmark_critical_forward_hooks", 

56 "benchmark_gated_hooks_fire", 

57 "benchmark_hook_functionality", 

58 "benchmark_hook_registry", 

59 # Gradient benchmarks 

60 "benchmark_backward_hooks", 

61 "benchmark_critical_backward_hooks", 

62 "benchmark_gradient_computation", 

63 # Generation benchmarks 

64 "benchmark_generation", 

65 "benchmark_generation_with_kv_cache", 

66 "benchmark_multiple_generation_calls", 

67 # Text quality benchmarks 

68 "benchmark_text_quality", 

69 # Weight processing benchmarks 

70 "benchmark_weight_modification", 

71 # Activation cache benchmarks 

72 "benchmark_activation_cache", 

73 "benchmark_run_with_cache", 

74]