replaced max with clamp

Author: mkhona-nvidia

emerging_optimizers/psgd/procrustes_step.py

@@ -43,7 +43,7 @@ def procrustes_step(Q: torch.Tensor, max_step_size: float = 0.125, eps: float =
     # Note: this function is written in fp32 to avoid numerical instability while computing the taylor expansion of the exponential map
     with utils.fp32_matmul_precision("highest"):
         R = Q.T - Q
-        R /= torch.max(norm_lower_bound_skew(R), eps)
+        R /= torch.clamp(norm_lower_bound_skew(R), min=eps)
         RQ = R @ Q
         # trace of RQ is always positive,
         # since tr(RQ) = ⟨R, Q⟩_F = ⟨Q^T - Q, Q⟩_F = ||Q||_F^2 - ⟨Q, Q⟩_F = ||Q||_F^2 - tr(Q^T Q) ≥ 0

emerging_optimizers/psgd/psgd_utils.py

@@ -85,7 +85,7 @@ def norm_lower_bound_spd(A: torch.Tensor, k: int = 4, half_iters: int = 2, eps:
     """
 
     # Compute scaling factor from the largest diagonal entry to prevent overflow/underflow
-    scale = torch.max(A.diagonal().amax(), eps)
+    scale = torch.clamp(A.diagonal().amax(), min=eps)
     A = A / scale
 
     bound_unnormalized = _subspace_iteration_bound(A, k=k, half_iters=half_iters, eps=eps)
@@ -113,7 +113,7 @@ def norm_lower_bound_skew(A: torch.Tensor, k: int = 32, half_iters: int = 2, eps
     """
 
     # Compute scaling factor from the max absolute value to prevent overflow/underflow
-    scale = torch.max(A.abs().amax(), eps)
+    scale = torch.clamp(A.abs().amax(), min=eps)
     A = A / scale
 
     bound_unnormalized = _subspace_iteration_bound(A, k=k, half_iters=half_iters, eps=eps)