Fix a bug with sparse updates in VariationalSGD optimizer.

PiperOrigin-RevId: 388353133

Author: jburnim

tensorflow_probability/python/optimizer/variational_sgd.py

@@ -251,8 +251,9 @@ def _resource_apply_sparse(self, grad, var, indices):
         self._burnin_max_learning_rate, self._max_learning_rate)
 
     learn_rate = tf.clip_by_value(
-        self._get_coordinatewise_learning_rate(grad, var), 0.,
-        tf.cast(max_learning_rate, var.dtype))
+        self._get_coordinatewise_learning_rate(
+            tf.IndexedSlices(grad, indices), var),
+        0., tf.cast(max_learning_rate, var.dtype))
     delta = grad * learn_rate
 
     return self._resource_scatter_add(var, indices, -delta)

tensorflow_probability/python/optimizer/variational_sgd_test.py

@@ -309,7 +309,6 @@ def testWithGlobalStep(self):
             init_step_value + 1, self.evaluate(sgd_optimizer.iterations))
 
   def testSparseBasic(self):
-    self.skipTest('b/195306553')
     for dtype in [tf.half, tf.float32, tf.float64]:
       with self.cached_session():
         var0 = tf.Variable([[1.1], [2.1]], dtype=dtype)