From a2fcd4dee5b2134a035d11ea2e4b3c7d339f50f3 Mon Sep 17 00:00:00 2001
From: Jiyan Yang <chocjy@fb.com>
Date: Wed, 28 Nov 2018 02:13:21 -0800
Subject: [PATCH] Ensure FP16 rowwise Adagrad can be run

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/12317

Reviewed By: hyuen

Differential Revision: D10190778

fbshipit-source-id: 720a9aaa4e6b1736023d8c6326a613e4ea592b31
---
 caffe2/python/optimizer.py      |  4 +++-
 caffe2/python/optimizer_test.py | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/caffe2/python/optimizer.py b/caffe2/python/optimizer.py
index ddd5871..b84f852 100644
--- a/caffe2/python/optimizer.py
+++ b/caffe2/python/optimizer.py
@@ -560,7 +560,7 @@ class AdagradOptimizer(Optimizer):
         )
 
         if self.rowWise:
-            assert self.engine == "SIMD", "Got {}".format(self.engine)
+            logger.info("Using engine {} for rowWise Adagrad".format(self.engine))
 
             shapes, types = workspace.InferShapesAndTypes([param_init_net])
             if str(param) not in shapes:
@@ -586,6 +586,8 @@ class AdagradOptimizer(Optimizer):
                     value=0.0
                 )
         else:
+            logger.info("Using engine {} for regular Adagrad".format(self.engine))
+
             if self.engine in FP16_ENGINES:
                 shapes, types = workspace.InferShapesAndTypes([param_init_net])
                 assert str(param) in shapes, shapes
diff --git a/caffe2/python/optimizer_test.py b/caffe2/python/optimizer_test.py
index bbb0c02..0dc8e55 100644
--- a/caffe2/python/optimizer_test.py
+++ b/caffe2/python/optimizer_test.py
@@ -137,6 +137,26 @@ class TestAdagrad(OptimizerTestBase, LRModificationTestBase, TestCase):
             workspace.FetchBlob(param)
 
 
+class TestRowWiseAdagrad(OptimizerTestBase, TestCase):
+    def build_optimizer(self, model, **kwargs):
+        self._skip_gpu = True
+        return build_adagrad(
+            model, base_learning_rate=1.0, lars=0.5, rowWise=True, **kwargs
+        )
+
+    def check_optimizer(self, optimizer):
+        self.assertFalse(optimizer.get_auxiliary_parameters().shared)
+        self.assertTrue(optimizer.get_auxiliary_parameters().local)
+        for param in optimizer.get_auxiliary_parameters().local:
+            workspace.FetchBlob(param)
+
+    def testDense(self):
+        raise unittest.SkipTest("no dense support")
+
+    def testGPUDense(self):
+        raise unittest.SkipTest("no dense support")
+
+
 class TestWngrad(OptimizerTestBase, LRModificationTestBase, TestCase):
     def build_optimizer(self, model, **kwargs):
         self._skip_gpu = True
-- 
2.7.4