Apply manual ruff fixes

k_diffusion/gns.py

@@ -5,8 +5,8 @@ class DDPGradientStatsHook:
     def __init__(self, ddp_module):
         try:
             ddp_module.register_comm_hook(self, self._hook_fn)
-        except AttributeError:
-            raise ValueError('DDPGradientStatsHook does not support non-DDP wrapped modules')
+        except AttributeError as ae:
+            raise ValueError('DDPGradientStatsHook does not support non-DDP wrapped modules') from ae
         self._clear_state()
 
     def _clear_state(self):

k_diffusion/models/axial_rope.py

@@ -84,7 +84,13 @@ def init(shape):
 
 
 class AxialRoPE(nn.Module):
-    def __init__(self, dim, n_heads, start_index=0, freqs_init=freqs_pixel_log(max_freq=10.0)):
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        start_index=0,
+        freqs_init=freqs_pixel_log(max_freq=10.0),  # noqa: B008
+    ):
         super().__init__()
         self.n_heads = n_heads
         self.start_index = start_index

k_diffusion/utils.py

@@ -180,8 +180,11 @@ def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., min_lr=0.,
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, "
+                "please use `get_last_lr()`.",
+                stacklevel=1,
+            )
 
         return self._get_closed_form_lr()
 
@@ -221,8 +224,11 @@ def __init__(self, optimizer, num_steps, decay=0.5, warmup=0., min_lr=0.,
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, "
+                "please use `get_last_lr()`.",
+                stacklevel=1,
+            )
 
         return self._get_closed_form_lr()
 
@@ -253,8 +259,11 @@ def __init__(self, optimizer, warmup=0., last_epoch=-1, verbose=False):
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, "
+                "please use `get_last_lr()`.",
+                stacklevel=1,
+            )
 
         return self._get_closed_form_lr()