Coverage for transformer_lens/conversion_utils/helpers/merge_quantiziation_fields.py: 93%

1"""Merge quantization fields helper. 

2 

3This module contains helper functions for merging quantization fields. 

4""" 

5 

6from typing import Any 

7 

8 

9def merge_quantization_fields(field_set: Any, quantization_fields: dict[str, Any]) -> Any: 

10 """Merge quantization fields into a field set. 

11 

12 Args: 

13 field_set: The field set to merge into. 

14 quantization_fields: The quantization fields to merge. 

15 

16 Returns: 

17 The merged field set (same object, modified in-place). 

18 """ 

19 for field_name, new_field_value in quantization_fields.items(): 

20 existing_field = field_set.fields.get(field_name) 

21 

22 # Check if existing field is None and raise error as expected by tests 

23 if existing_field is None: 

24 raise RuntimeError( 

25 "Attempted to merge quantization field into existing conversion without original field configured" 

26 ) 

27 

28 # Handle different cases based on the types of existing and new fields 

29 if isinstance(new_field_value, tuple) and len(new_field_value) == 2: 

30 # new_field_value is (str, TensorConversionSet) 

31 new_remote, new_sub_wcs = new_field_value 

32 

33 if isinstance(existing_field, tuple) and len(existing_field) == 2: 

34 # existing_field is also (str, TensorConversionSet) 

35 existing_remote, existing_sub_wcs = existing_field 

36 

37 # Check if the second element is a TensorConversionSet-like object 

38 if hasattr(existing_sub_wcs, "fields") and hasattr(new_sub_wcs, "fields"): 38 ↛ 43line 38 didn't jump to line 43 because the condition on line 38 was always true

39 merge_quantization_fields(existing_sub_wcs, new_sub_wcs.fields) 

40 # Update the remote field name 

41 field_set.fields[field_name] = (new_remote, existing_sub_wcs) 

42 else: 

43 raise RuntimeError( 

44 "Attempted to merge TensorConversionSet into a field that is not configured as a TensorConversionSet" 

45 ) 

46 else: 

47 # existing_field is not a tuple, but new_field_value is 

48 raise RuntimeError( 

49 "Attempted to merge TensorConversionSet into a field that is not configured as a TensorConversionSet" 

50 ) 

51 else: 

52 # new_field_value is a simple value (like torch.Tensor) 

53 field_set.fields[field_name] = new_field_value 

54 

55 return field_set