Lint and format changes

Author: ParthS007

opacus/optimizers/__init__.py

@@ -15,21 +15,21 @@
 from .adaclipoptimizer import AdaClipDPOptimizer
 from .ddp_perlayeroptimizer import SimpleDistributedPerLayerOptimizer
 from .ddpoptimizer import DistributedDPOptimizer
+from .ddpoptimizer_automatic_clipping import (
+    DistributedDPAutomaticClippingOptimizer,
+    DistributedDPPerLayerAutomaticClippingOptimizer,
+)
 from .ddpoptimizer_fast_gradient_clipping import (
     DistributedDPOptimizerFastGradientClipping,
 )
 from .fsdpoptimizer_fast_gradient_clipping import FSDPOptimizerFastGradientClipping
 from .optimizer import DPOptimizer
-from .optimizer_fast_gradient_clipping import DPOptimizerFastGradientClipping
-from .perlayeroptimizer import DPPerLayerOptimizer
 from .optimizer_automatic_clipping import (
     DPAutomaticClippingOptimizer,
     DPPerLayerAutomaticClippingOptimizer,
 )
-from .ddpoptimizer_automatic_clipping import (
-    DistributedDPAutomaticClippingOptimizer,
-    DistributedDPPerLayerAutomaticClippingOptimizer,
-)
+from .optimizer_fast_gradient_clipping import DPOptimizerFastGradientClipping
+from .perlayeroptimizer import DPPerLayerOptimizer
 
 
 __all__ = [

opacus/optimizers/ddpoptimizer_automatic_clipping.py

@@ -17,12 +17,11 @@
 from typing import Callable, List, Optional
 
 import torch
-from torch.optim import Optimizer
-
 from opacus.optimizers.optimizer_automatic_clipping import (
     DPAutomaticClippingOptimizer,
     DPPerLayerAutomaticClippingOptimizer,
 )
+from torch.optim import Optimizer
 
 
 class DistributedDPAutomaticClippingOptimizer(DPAutomaticClippingOptimizer):

opacus/optimizers/optimizer_automatic_clipping.py

@@ -18,9 +18,12 @@
 from typing import List
 
 import torch
-from opacus.optimizers.optimizer import _check_processed_flag, _mark_as_processed
+from opacus.optimizers.optimizer import (
+    DPOptimizer,
+    _check_processed_flag,
+    _mark_as_processed,
+)
 from opacus.optimizers.perlayeroptimizer import DPPerLayerOptimizer
-from opacus.optimizers.optimizer import DPOptimizer
 
 
 class DPAutomaticClippingOptimizer(DPOptimizer):
@@ -102,8 +105,9 @@ class DPPerLayerAutomaticClippingOptimizer(DPPerLayerOptimizer):
     Per-layer variant of automatic clipping.
 
     For each parameter (layer), we compute the per-sample clip factor using the
-    corresponding per-layer ``max_grad_norm``:
-        ``per_sample_clip_factor = max_grad_norm / (per_sample_norms + 0.01)``
+    corresponding per-layer ``max_grad_norm``::
+
+        per_sample_clip_factor = max_grad_norm / (per_sample_norms + 0.01)
 
     This allows each layer to have different clipping behavior based on its own
     gradient magnitude distribution, which can improve training stability and

opacus/tests/multigpu_automatic_clipping_test.py

@@ -115,7 +115,9 @@ def demo_basic(rank, weight, world_size, dp, clipping):
         if clipping == "automatic":
             assert isinstance(optimizer, DistributedDPAutomaticClippingOptimizer)
         elif clipping == "automatic_per_layer":
-            assert isinstance(optimizer, DistributedDPPerLayerAutomaticClippingOptimizer)
+            assert isinstance(
+                optimizer, DistributedDPPerLayerAutomaticClippingOptimizer
+            )
 
     for x, y in data_loader:
         outputs = ddp_model(x.to(rank))
@@ -167,9 +169,7 @@ def test_automatic_clipping_gradient_correct(self) -> None:
             clipping=None,
         )
 
-        self.assertTrue(
-            torch.allclose(weight_dp, weight_nodp, atol=1e-5, rtol=1e-3)
-        )
+        self.assertTrue(torch.allclose(weight_dp, weight_nodp, atol=1e-5, rtol=1e-3))
 
     @unittest.skipIf(torch.cuda.device_count() < 2, "Need at least 2 GPUs")
     def test_automatic_per_layer_clipping_gradient_correct(self) -> None:
@@ -199,9 +199,7 @@ def test_automatic_per_layer_clipping_gradient_correct(self) -> None:
             clipping=None,
         )
 
-        self.assertTrue(
-            torch.allclose(weight_dp, weight_nodp, atol=1e-5, rtol=1e-3)
-        )
+        self.assertTrue(torch.allclose(weight_dp, weight_nodp, atol=1e-5, rtol=1e-3))
 
     @unittest.skipIf(torch.cuda.device_count() < 2, "Need at least 2 GPUs")
     def test_automatic_clipping_optimizer_type(self) -> None: