feature: copy_stochastic

Author: kozistr

pytorch_optimizer/__init__.py

@@ -168,6 +168,7 @@
 from pytorch_optimizer.optimizer.utils import (
     CPUOffloadOptimizer,
     clip_grad_norm,
+    copy_stochastic,
     disable_running_stats,
     enable_running_stats,
     get_global_gradient_norm,

pytorch_optimizer/optimizer/utils.py

@@ -310,3 +310,26 @@ def reg_noise(
         loss.add_(reg - noise.mul_(noise_coef).sum())
 
     return loss
+
+
+@torch.no_grad()
+def copy_stochastic(target: torch.Tensor, source: torch.Tensor) -> None:
+    r"""Copy stochastic.
+
+    reference: https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905
+
+    :param target: torch.Tensor. bfloat16 tensor.
+    :param source: torch.Tensor. float32 tensor.
+    """
+    result = torch.randint_like(
+        source,
+        dtype=torch.int32,
+        low=0,
+        high=1 << 16,
+    )
+
+    result.add_(source.view(dtype=torch.int32))
+
+    result.bitwise_and_(-65536)
+
+    target.copy_(result.view(dtype=torch.float32))