Revert "[BE] remove unnecessary _dispatch_sqrt by using ** 0.5 (pytorch#131358)"

This reverts commit 4c7f22d.

Reverted pytorch#131358 on behalf of https://github.com/janeyx99 due to Internal uses this private API and landing that has been a pain so we're reverting this first ([comment](pytorch#131358 (comment)))

Author: pytorchmergebot

torch/optim/adam.py

@@ -10,6 +10,7 @@
     _default_to_fused_or_foreach,
     _differentiable_doc,
     _disable_dynamo_if_unsupported,
+    _dispatch_sqrt,
     _foreach_doc,
     _fused_doc,
     _get_capturable_supported_devices,
@@ -424,7 +425,7 @@ def _single_tensor_adam(
 
             step_size = lr / bias_correction1
 
-            bias_correction2_sqrt = bias_correction2**0.5
+            bias_correction2_sqrt = _dispatch_sqrt(bias_correction2)
 
             if amsgrad:
                 # Maintains the maximum of all 2nd moment running avg. till now
@@ -595,7 +596,7 @@ def _multi_tensor_adam(
 
             step_size = _stack_if_compiling([(lr / bc) * -1 for bc in bias_correction1])
 
-            bias_correction2_sqrt = [bc**0.5 for bc in bias_correction2]  # type: ignore[arg-type]
+            bias_correction2_sqrt = [_dispatch_sqrt(bc) for bc in bias_correction2]  # type: ignore[arg-type]
 
             if amsgrad:
                 # Maintains the maximum of all 2nd moment running avg. till now

torch/optim/adamw.py

@@ -10,6 +10,7 @@
     _default_to_fused_or_foreach,
     _differentiable_doc,
     _disable_dynamo_if_unsupported,
+    _dispatch_sqrt,
     _foreach_doc,
     _fused_doc,
     _get_capturable_supported_devices,
@@ -425,7 +426,7 @@ def _single_tensor_adamw(
 
             step_size = lr / bias_correction1
 
-            bias_correction2_sqrt = bias_correction2**0.5
+            bias_correction2_sqrt = _dispatch_sqrt(bias_correction2)
 
             if amsgrad:
                 # Maintains the maximum of all 2nd moment running avg. till now
@@ -592,7 +593,7 @@ def _multi_tensor_adamw(
             step_size = _stack_if_compiling([(lr / bc) * -1 for bc in bias_correction1])
 
             bias_correction2_sqrt = [
-                bc**0.5 for bc in bias_correction2  # type: ignore[arg-type]
+                _dispatch_sqrt(bc) for bc in bias_correction2  # type: ignore[arg-type]
             ]
 
             if amsgrad:

torch/optim/nadam.py

@@ -10,6 +10,7 @@
     _default_to_fused_or_foreach,
     _differentiable_doc,
     _disable_dynamo_if_unsupported,
+    _dispatch_sqrt,
     _foreach_doc,
     _get_capturable_supported_devices,
     _get_scalar_dtype,
@@ -488,7 +489,8 @@ def _multi_tensor_nadam(
             torch._foreach_sqrt_(bias_correction_sqrt)
         else:
             bias_correction_sqrt = [
-                (1 - beta2 ** _get_value(step)) ** 0.5 for step in grouped_state_steps
+                _dispatch_sqrt(1 - beta2 ** _get_value(step))
+                for step in grouped_state_steps
             ]
             mus = [
                 beta1 * (1.0 - 0.5 * (0.96 ** (_get_value(step) * momentum_decay)))

torch/optim/optimizer.py

@@ -2,6 +2,7 @@
 # mypy: allow-untyped-defs
 """Base optimizer."""
 import functools
+import math
 import warnings
 from collections import defaultdict, OrderedDict
 from copy import deepcopy
@@ -112,6 +113,15 @@ def _stack_if_compiling(x):
         return x
 
 
+def _dispatch_sqrt(
+    x: float,
+):  # float annotation is needed because of torchscript type inference
+    if not torch.jit.is_scripting() and isinstance(x, torch.Tensor):
+        return x.sqrt()
+    else:
+        return math.sqrt(x)
+
+
 def _disable_dynamo_if_unsupported(single_tensor_fn=None):
     # workaround for torchscript BC
     # it requires all called functions to be in the

torch/optim/radam.py

@@ -11,6 +11,7 @@
     _default_to_fused_or_foreach,
     _differentiable_doc,
     _disable_dynamo_if_unsupported,
+    _dispatch_sqrt,
     _foreach_doc,
     _get_capturable_supported_devices,
     _get_scalar_dtype,
@@ -504,13 +505,12 @@ def _multi_tensor_radam(
             del bias_correction1
         else:
             rect = [
-                (
+                _dispatch_sqrt(
                     (rho_t - 4)  # type: ignore[arg-type]
                     * (rho_t - 2)
                     * rho_inf
                     / ((rho_inf - 4) * (rho_inf - 2) * rho_t)
                 )
-                ** 0.5
                 if rho_t > 5
                 else 0
                 for rho_t in rho_t_list
@@ -524,7 +524,7 @@ def _multi_tensor_radam(
                 (lr * rect / bc) * -1 for rect, bc in zip(unrectified, bias_correction1)
             ]
             bias_correction2 = [
-                ((1 - beta2 ** _get_value(step)) ** 0.5) * (lr * rect / bc) * -1
+                _dispatch_sqrt(1 - beta2 ** _get_value(step)) * (lr * rect / bc) * -1
                 for step, rect, bc in zip(grouped_state_steps, rect, bias_correction1)
             ]