added formatting checks, license, etc

Signed-off-by: mikail <mkhona@nvidia.com>

Author: mkhona-nvidia

emerging_optimizers/scalar_optimizers/__init__.py

@@ -13,10 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from .adam import calculate_adam_update
-from .ademamix import calculate_sim_ademamix_update, calculate_ademamix_update
-from .signum import calculate_signum_update
+from .ademamix import calculate_ademamix_update, calculate_sim_ademamix_update
 from .laprop import calculate_laprop_update
 from .lion import calculate_lion_update
+from .signum import calculate_signum_update
+
 
 __all__ = [
     "calculate_adam_update",

emerging_optimizers/scalar_optimizers/adam.py

@@ -12,9 +12,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import torch
 from typing import Tuple
 
+import torch
+
+
 __all__ = [
     "calculate_adam_update",
 ]

emerging_optimizers/scalar_optimizers/ademamix.py

@@ -12,9 +12,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import torch
-from typing import Tuple, Optional
 import math
+from typing import Optional, Tuple
+
+import torch
+
 
 __all__ = [
     "calculate_sim_ademamix_update",

emerging_optimizers/scalar_optimizers/laprop.py

@@ -12,9 +12,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import torch
 from typing import Tuple
 
+import torch
+
 
 __all__ = [
     "calculate_laprop_update",

emerging_optimizers/scalar_optimizers/lion.py

@@ -12,9 +12,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import torch
 from typing import Optional
 
+import torch
+
+
 __all__ = [
     "calculate_lion_update",
 ]

emerging_optimizers/scalar_optimizers/signum.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 import torch
 
+
 __all__ = [
     "calculate_signum_update",
 ]
@@ -34,7 +35,7 @@ def calculate_signum_update(
 
     This function performs the computation of 1 step of sign-SGD or Signum.
     Based on https://arxiv.org/abs/1802.04434.
-    When using signSGD with shape scaling, general recommendation is to 
+    When using signSGD with shape scaling, general recommendation is to
     scale :math:`lr = \\text{adam lr} \\cdot \\text{network width} \\cdot \\frac{2}{\\text{rows} + \\text{cols}}`.
     This is for learning rate transfer with width scaling (https://arxiv.org/abs/2506.07254v1).
 

emerging_optimizers/soap/soap.py

@@ -13,20 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from itertools import chain
-from typing import Iterable, Callable, Tuple, List, Optional, Union, override
+from typing import Callable, Iterable, List, Optional, Tuple, Union, override
 
 import torch
 import torch.optim as optim
-
 from absl import logging
 
+from emerging_optimizers import utils
+from emerging_optimizers.scalar_optimizers import calculate_adam_update
 from emerging_optimizers.soap.soap_utils import (
     get_eigenbasis_eigh,
     get_eigenbasis_qr,
 )
 
-from emerging_optimizers import utils
-from emerging_optimizers.scalar_optimizers import calculate_adam_update
 
 __all__ = [
     "SOAP",
@@ -104,11 +103,10 @@ def __init__(
         power_iter_steps: int = 1,
         max_update_rms: float = 0.0,
     ) -> None:
-
         # Check for betas.
         if betas is None:
             betas = (0.95, 0.95)
-            logging.debug("betas not provided. Setting betas equal to " f"betas = {betas} by default.")
+            logging.debug(f"betas not provided. Setting betas equal to betas = {betas} by default.")
 
         # Check for update criteria
         if use_adaptive_criteria:

emerging_optimizers/soap/soap_utils.py

@@ -13,10 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import List, Optional, Tuple
+
 import torch
 
 from emerging_optimizers import utils
 
+
 __all__ = [
     "get_eigenbasis_eigh",
     "get_eigenbasis_qr",

emerging_optimizers/utils/eig.py

@@ -12,14 +12,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from absl import logging
 from typing import Optional
 
 import torch
+from absl import logging
 from torch import Tensor
 
 from emerging_optimizers import utils
 
+
 __all__ = ["eigh_with_fallback", "eig_orthogonal_iteration", "adaptive_early_exit_criteria"]
 
 

emerging_optimizers/utils/precondition_schedules.py

@@ -16,6 +16,7 @@
 from abc import ABC, abstractmethod
 from typing import Dict
 
+
 __all__ = [
     "LinearSchedule",
     "CosineSchedule",