added cpu and gpu flag

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

Author: mkhona-nvidia

emerging_optimizers/scalar_optimizers/laprop.py

@@ -76,7 +76,7 @@ def calculate_laprop_update(
 
     # construct the denominator of the inner ADAM optimizer
     second_moment = exp_avg_sq / bias_correction2
-    second_moment = second_moment.sqrt() + eps
+    second_moment.sqrt_().add_(eps)
 
     normalized_grad = grad / second_moment
 

tests/ci/L0_Tests_CPU.sh

@@ -15,3 +15,5 @@ export TORCH_COMPILE_DISABLE=1
 set -o pipefail
 torchrun --nproc_per_node=8 --no-python coverage run -p tests/test_distributed_muon_utils_cpu.py
 torchrun --nproc_per_node=4 --no-python coverage run -p tests/test_distributed_muon_utils_cpu.py
+coverage run -p --source=emerging_optimizers tests/test_scalar_optimizers.py --device=cpu 
+

tests/ci/L0_Tests_GPU.sh

@@ -20,4 +20,5 @@ coverage run -p --source=emerging_optimizers tests/test_soap_functions.py
 coverage run -p --source=emerging_optimizers tests/test_soap_utils.py
 coverage run -p --source=emerging_optimizers tests/soap_smoke_test.py
 coverage run -p --source=emerging_optimizers tests/soap_mnist_test.py
-coverage run -p --source=emerging_optimizers tests/test_scalar_optimizers.py
+coverage run -p --source=emerging_optimizers tests/test_scalar_optimizers.py --device=auto
+

tests/ci/L1_Tests_GPU.sh

@@ -18,4 +18,4 @@ python tests/test_orthogonalized_optimizer.py
 python tests/test_soap_functions.py
 python tests/test_soap_utils.py
 python tests/soap_smoke_test.py
-python tests/test_scalar_optimizers.py
+python tests/test_scalar_optimizers.py --device=auto

tests/test_scalar_optimizers.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch
+from absl import flags
 from absl.testing import absltest, parameterized
 
 from emerging_optimizers.scalar_optimizers import (
@@ -23,22 +24,54 @@
 )
 
 
+# Define command line flags
+flags.DEFINE_string("device", "cpu", "Device to run tests on: 'cpu', 'cuda', or 'auto'")
+flags.DEFINE_integer("seed", 42, "Random seed for reproducible tests")
+flags.DEFINE_boolean("skip_gpu_tests", False, "Skip GPU tests even if CUDA is available")
+
+FLAGS = flags.FLAGS
+
+
 # Base class for tests requiring seeding for determinism
 class BaseTestCase(parameterized.TestCase):
     def setUp(self):
-        """Set random seed before each test."""
-        # Set seed for PyTorch
-        torch.manual_seed(42)
+        """Set random seed and device before each test."""
+        # Set seed for PyTorch (using seed from flags)
+        torch.manual_seed(FLAGS.seed)
         # Set seed for CUDA if available
         if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(42)
+            torch.cuda.manual_seed_all(FLAGS.seed)
+
+        # Set up device based on flags
+        self.device = self._get_test_device()
+
+    def _get_test_device(self):
+        """Get the device to use for testing based on flags."""
+        if FLAGS.device == "auto":
+            return "cuda" if torch.cuda.is_available() and not FLAGS.skip_gpu_tests else "cpu"
+        elif FLAGS.device == "cuda":
+            if not torch.cuda.is_available():
+                self.skipTest("CUDA not available")
+            if FLAGS.skip_gpu_tests:
+                self.skipTest("GPU tests skipped by flag")
+            return "cuda"
+        else:
+            return "cpu"
+
+    def _move_to_device(self, *tensors):
+        """Helper method to move tensors to the test device."""
+        return tuple(tensor.to(self.device) for tensor in tensors)
 
 
 class ScalarOptimizerTest(BaseTestCase):
     def test_calculate_adam_update_simple(self) -> None:
         exp_avg_initial = torch.tensor([[1.0]])
         exp_avg_sq_initial = torch.tensor([[2.0]])
         grad = torch.tensor([[0.5]])
+
+        # Move tensors to the test device
+        exp_avg_initial, exp_avg_sq_initial, grad = self._move_to_device(exp_avg_initial, exp_avg_sq_initial, grad)
+
         betas = (0.9, 0.99)
         eps = 1e-8
         step = 10
@@ -59,7 +92,7 @@ def test_calculate_adam_update_simple(self) -> None:
             eps=eps,
         )
 
-        initial_param_val_tensor = torch.tensor([[10.0]])
+        initial_param_val_tensor = torch.tensor([[10.0]]).to(self.device)
         param = torch.nn.Parameter(initial_param_val_tensor.clone())
         param.grad = grad.clone()
 
@@ -249,6 +282,23 @@ def test_calculate_sim_ademamix_update_with_zero_momentum_and_alpha_equals_rmspr
         expected_param_val_after_step = initial_param_val_tensor - lr * sim_ademamix_update
         torch.testing.assert_close(param.data, expected_param_val_after_step, atol=1e-6, rtol=1e-6)
 
+    def test_device_functionality(self) -> None:
+        """Test that tensors are correctly moved to the specified device."""
+        # Create test tensors
+        tensor1 = torch.tensor([1.0, 2.0, 3.0])
+        tensor2 = torch.tensor([[1.0], [2.0]])
+
+        # Move to test device
+        tensor1_device, tensor2_device = self._move_to_device(tensor1, tensor2)
+
+        # Verify they are on the correct device
+        self.assertEqual(str(tensor1_device.device), self.device)
+        self.assertEqual(str(tensor2_device.device), self.device)
+
+        # Verify values are preserved
+        torch.testing.assert_close(tensor1_device.cpu(), tensor1, atol=1e-6, rtol=1e-6)
+        torch.testing.assert_close(tensor2_device.cpu(), tensor2, atol=1e-6, rtol=1e-6)
+
 
 if __name__ == "__main__":
     absltest.main()