55
66
77def concat_vectors_and_pad (vec_list , max_ ):
8+ """
9+ Concatenates a list of input vectors and pads them to match a specified maximum
10+ length.
11+
12+ This function takes a list of input vectors, concatenates them along a specified
13+ dimension (dim=0), and then pads the concatenated vector to achieve a specified
14+ maximum length. The padding is done with zeros.
15+
16+ Args:
17+ vec_list (list of torch.Tensor): List of input vectors to concatenate and pad.
18+ max_ (int): The maximum length of the concatenated and padded vector.
19+
20+ Returns:
21+ torch.Tensor: The concatenated and padded vector.
22+
23+ Raises:
24+ AssertionError: If the length of 'vec_list' is not greater than 0, or if it
25+ exceeds 'max_len', or if 'max_len' is not greater than 0.
26+
27+ Example:
28+ >>> input_tensors = [torch.tensor([1, 2]), torch.tensor([3, 4, 5])]
29+ >>> max_length = 5
30+ >>> concatenated_padded = concat_vectors_and_pad(input_tensors, max_length)
31+ >>> print(concatenated_padded)
32+ tensor([1, 2, 3, 4, 5])
33+ """
834 assert len (vec_list ) > 0
935 assert len (vec_list ) <= max_
1036 assert max_ > 0
@@ -27,10 +53,29 @@ class LightwoodAutocast:
2753 """
2854 Equivalent to torch.cuda.amp.autocast, but checks device compute capability
2955 to activate the feature only when the GPU has tensor cores to leverage AMP.
56+
57+ **Attributes:**
58+
59+ * `active` (bool): Whether AMP is currently active. This attribute is at the class
60+ level
61+
62+ **Usage:**
63+
64+ ```python
65+ >>> import lightwood.helpers.torch as lt
66+ >>> with lt.LightwoodAutocast():
67+ ... # This code will be executed in AMP mode.
68+ ... pass
3069 """
3170 active = False
3271
3372 def __init__ (self , enabled = True ):
73+ """
74+ Initializes the context manager for Automatic Mixed Precision (AMP) functionality.
75+
76+ Args:
77+ enabled (bool, optional): Whether to enable AMP. Defaults to True.
78+ """
3479 self .major = 0 # GPU major version
3580 torch_version = [int (i ) for i in torch .__version__ .split ('.' )[:- 1 ]]
3681
@@ -50,12 +95,18 @@ def __init__(self, enabled=True):
5095 LightwoodAutocast .active = self ._enabled
5196
5297 def __enter__ (self ):
98+ """
99+ * `__enter__()`: Enters the context manager and enables AMP if it is not already enabled.
100+ """
53101 if self ._enabled :
54102 self .prev = torch .is_autocast_enabled ()
55103 torch .set_autocast_enabled (self ._enabled )
56104 torch .autocast_increment_nesting ()
57105
58106 def __exit__ (self , * args ):
107+ """
108+ * `__exit__()`: Exits the context manager and disables AMP.
109+ """
59110 if self ._enabled :
60111 # Drop the cache when we exit to a nesting level that's outside any instance of autocast
61112 if torch .autocast_decrement_nesting () == 0 :
@@ -64,6 +115,9 @@ def __exit__(self, *args):
64115 return False
65116
66117 def __call__ (self , func ):
118+ """
119+ * `__call__(self, func)`: Returns a decorated function that enables AMP when it is called.
120+ """
67121 @functools .wraps (func )
68122 def decorate_autocast (* args , ** kwargs ):
69123 with self :
0 commit comments