Repository: tensorflow/addons

Issue:
SGDW with weight decay schedule fails to run on GPU
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): `Deep Learning Linux VM (Debian Stretch) GCE Image`
- TensorFlow version and how it was installed (source or binary): `Binary TF_v2.3.0rc2`
- TensorFlow-Addons version and how it was installed (source or binary): `Binary tfa-nightly_v0.11.0.dev20200723223438`
- Python version: `3.7.6`
- Is GPU used? (yes/no): `yes`

**Describe the bug**

#2029 adds support for weighted decay schedule for `DecoupledWeightDecayExtension` i.e. for SGDW, AdamW, etc. Weight decay scheduling is now supported using `tf.keras.optimizers.schedules.LearningRateScheduler`, but the same fails to work when executed on a GPU. (when executed on a CPU, it runs fine as is!)

**Code to reproduce the issue**

```python
import tensorflow as tf
import tensorflow_addons as tfa

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(2, 2, padding='same', input_shape=(5, 5, 1)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(2, input_shape=(10,))
])
model.summary()

init_lr = 0.1
decay = 0.2
boundaries = [5, 10, 15]
lr_boundaries = [(2 * epoch) for epoch in sorted(boundaries)]
lr_values = [init_lr * (decay ** idx) for idx in range(len(lr_boundaries) + 1)]
wd = 0.0005
wd_values = [(wd * lrx) for lrx in lr_values]

lr_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(lr_boundaries, lr_values)
wd_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(lr_boundaries, wd_values)

opt = tfa.optimizers.SGDW(wd_schedule, lr_schedule, momentum=0.9, nesterov=True)
model.compile(loss='sparse_categorical_crossentropy', optimizer=opt)

fake_x, fake_y = tf.random.normal([32, 5, 5, 1]), tf.random.uniform([32], 0, 2)
model.fit(fake_x, fake_y, epochs=20, batch_size=16)
```

**Other info / logs**

```
2020-07-27 20:20:43.167840: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
/opt/conda/lib/python3.7/site-packages/tensorflow_addons/utils/ensure_tf_install.py:68: UserWarning: Tensorflow Addons supports using Python ops for all Tensorflow versions above or equal to 2.2.0 and strictly below 2.3.0 (nightly versions are not supported). 
 The versions of TensorFlow you are currently using is 2.3.0-rc2 and is not supported. 
Some things might work, some things might not.
If you were to encounter a bug, do not file an issue.
If you want to make sure you're using a tested and supported configuration, either change the TensorFlow version or the TensorFlow Addons's version. 
You can find the compatibility matrix in TensorFlow Addon's readme:
https://github.com/tensorflow/addons
  UserWarning,
2020-07-27 20:20:44.470285: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcuda.so.1
2020-07-27 20:20:45.055810: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.056455: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 0 with properties: 
pciBusID: 0000:00:04.0 name: Tesla T4 computeCapability: 7.5
coreClock: 1.59GHz coreCount: 40 deviceMemorySize: 14.73GiB deviceMemoryBandwidth: 298.08GiB/s
2020-07-27 20:20:45.056514: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
2020-07-27 20:20:45.058582: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcublas.so.10
2020-07-27 20:20:45.060424: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcufft.so.10
2020-07-27 20:20:45.060759: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcurand.so.10
2020-07-27 20:20:45.062699: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusolver.so.10
2020-07-27 20:20:45.063780: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusparse.so.10
2020-07-27 20:20:45.067730: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudnn.so.7
2020-07-27 20:20:45.067894: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.068645: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.069201: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1858] Adding visible gpu devices: 0
2020-07-27 20:20:45.069573: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN)to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2020-07-27 20:20:45.078074: I tensorflow/core/platform/profile_utils/cpu_utils.cc:104] CPU Frequency: 2300000000 Hz
2020-07-27 20:20:45.078334: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x55c5b4017cb0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-07-27 20:20:45.078440: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2020-07-27 20:20:45.176377: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.177148: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x55c5b402b9e0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2020-07-27 20:20:45.177183: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Tesla T4, Compute Capability 7.5
2020-07-27 20:20:45.177413: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.177980: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 0 with properties: 
pciBusID: 0000:00:04.0 name: Tesla T4 computeCapability: 7.5
coreClock: 1.59GHz coreCount: 40 deviceMemorySize: 14.73GiB deviceMemoryBandwidth: 298.08GiB/s
2020-07-27 20:20:45.178028: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
2020-07-27 20:20:45.178081: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcublas.so.10
2020-07-27 20:20:45.178104: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcufft.so.10
2020-07-27 20:20:45.178125: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcurand.so.10
2020-07-27 20:20:45.178144: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusolver.so.10
2020-07-27 20:20:45.178163: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusparse.so.10
2020-07-27 20:20:45.178183: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudnn.so.7
2020-07-27 20:20:45.178266: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.178859: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.179375: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1858] Adding visible gpu devices: 0
2020-07-27 20:20:45.179422: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
2020-07-27 20:20:45.877424: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1257] Device interconnect StreamExecutor with strength 1 edge matrix:
2020-07-27 20:20:45.877479: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1263]      0 
2020-07-27 20:20:45.877490: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 0:   N 
2020-07-27 20:20:45.877756: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.878463: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-07-27 20:20:45.879064: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1402] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 13936 MB memory) -> physical GPU (device: 0, name: Tesla T4, pci bus id: 0000:00:04.0, compute capability: 7.5)
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d (Conv2D)              (None, 5, 5, 2)           10        
_________________________________________________________________
batch_normalization (BatchNo (None, 5, 5, 2)           8         
_________________________________________________________________
global_average_pooling2d (Gl (None, 2)                 0         
_________________________________________________________________
dense (Dense)                (None, 2)                 6         
=================================================================
Total params: 24
Trainable params: 20
Non-trainable params: 4
_________________________________________________________________
Epoch 1/20
Traceback (most recent call last):
  File "tfa-issue.py", line 27, in <module>
    model.fit(fake_x, fake_y, epochs=20, batch_size=16)
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 108, in _method_wrapper
    return method(self, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 1098, in fit
    tmp_logs = train_function(iterator)
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 780, in __call__
    result = self._call(*args, **kwds)
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 840, in _call
    return self._stateless_fn(*args, **kwds)
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2829, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1848, in _filtered_call
    cancellation_manager=cancellation_manager)
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1924, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 550, in call
    ctx=ctx)
  File "/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation sequential/conv2d/Conv2D/ReadVariableOp: Could not satisfy explicit device specification '' because the node {{colocation_node sequential/conv2d/Conv2D/ReadVariableOp}} was colocated with a group of nodes that required incompatible device '/job:localhost/replica:0/task:0/device:GPU:0'. All available devices [/job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:XLA_CPU:0, /job:localhost/replica:0/task:0/device:XLA_GPU:0, /job:localhost/replica:0/task:0/device:GPU:0]. 
Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=2 requested_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' assigned_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' resource_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' supported_device_types_=[CPU] possible_devices_=[]
AssignSubVariableOp: GPU CPU XLA_CPU XLA_GPU 
Mul: GPU CPU XLA_CPU XLA_GPU 
Sum: GPU CPU XLA_CPU XLA_GPU 
ResourceApplyKerasMomentum: GPU CPU XLA_CPU XLA_GPU 
LogicalAnd: GPU CPU XLA_CPU XLA_GPU 
Identity: GPU CPU XLA_CPU XLA_GPU 
LessEqual: GPU CPU XLA_CPU XLA_GPU 
Pack: GPU CPU XLA_CPU XLA_GPU 
Const: CPU 
Cast: GPU CPU XLA_CPU XLA_GPU 
ReadVariableOp: GPU CPU XLA_CPU XLA_GPU 
_Arg: GPU CPU XLA_CPU XLA_GPU 
Greater: GPU CPU XLA_CPU XLA_GPU 

Colocation members, user-requested devices, and framework assigned devices, if any:
  sequential_conv2d_conv2d_readvariableop_resource (_Arg)  framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
  sgdw_sgdw_update_resourceapplykerasmomentum_accum (_Arg)  framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
  sequential/conv2d/Conv2D/ReadVariableOp (ReadVariableOp) 
  SGDW/SGDW/update/PiecewiseConstant/Const (Const) 
  SGDW/SGDW/update/PiecewiseConstant/Const_1 (Const) 
  SGDW/SGDW/update/PiecewiseConstant/Const_2 (Const) 
  SGDW/SGDW/update/PiecewiseConstant/Const_3 (Const) 
  SGDW/SGDW/update/PiecewiseConstant/Const_4 (Const) 
  SGDW/SGDW/update/PiecewiseConstant/Const_5 (Const) 
  SGDW/SGDW/update/PiecewiseConstant/Const_6 (Const) 
  SGDW/SGDW/update/PiecewiseConstant/Cast (Cast) 
  SGDW/SGDW/update/PiecewiseConstant/Cast_1 (Cast) 
  SGDW/SGDW/update/PiecewiseConstant/Cast_2 (Cast) 
  SGDW/SGDW/update/PiecewiseConstant/LessEqual (LessEqual) 
  SGDW/SGDW/update/PiecewiseConstant/Greater (Greater) 
  SGDW/SGDW/update/PiecewiseConstant/Greater_1 (Greater) 
  SGDW/SGDW/update/PiecewiseConstant/LessEqual_1 (LessEqual) 
  SGDW/SGDW/update/PiecewiseConstant/and (LogicalAnd) 
  SGDW/SGDW/update/PiecewiseConstant/Greater_2 (Greater) 
  SGDW/SGDW/update/PiecewiseConstant/LessEqual_2 (LessEqual) 
  SGDW/SGDW/update/PiecewiseConstant/and_1 (LogicalAnd) 
  SGDW/SGDW/update/PiecewiseConstant/case/preds_c (Pack) 
  SGDW/SGDW/update/PiecewiseConstant/case/Cast (Cast) 
  SGDW/SGDW/update/PiecewiseConstant/case/Const (Const) 
  SGDW/SGDW/update/PiecewiseConstant/case/num_true_conds (Sum) 
  SGDW/SGDW/update/PiecewiseConstant/case/n_true_conds (Const) 
  SGDW/SGDW/update/PiecewiseConstant/case/LessEqual (LessEqual) 
  SGDW/SGDW/update/PiecewiseConstant/case/Assert/Const (Const) 
  SGDW/SGDW/update/PiecewiseConstant/case/Assert/AssertGuard/Identity (Identity) 
  SGDW/SGDW/update/PiecewiseConstant/case/cond/Identity (Identity) 
  SGDW/SGDW/update/mul/ReadVariableOp (ReadVariableOp) 
  SGDW/SGDW/update/mul (Mul) 
  SGDW/SGDW/update/AssignSubVariableOp (AssignSubVariableOp) 
  SGDW/SGDW/update/Cast (Cast) 
  SGDW/SGDW/update/PiecewiseConstant_1/Const (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/Const_1 (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/Const_2 (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/Const_3 (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/Const_4 (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/Const_5 (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/Const_6 (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/Cast (Cast) 
  SGDW/SGDW/update/PiecewiseConstant_1/Cast_1 (Cast) 
  SGDW/SGDW/update/PiecewiseConstant_1/Cast_2 (Cast) 
  SGDW/SGDW/update/PiecewiseConstant_1/LessEqual (LessEqual) 
  SGDW/SGDW/update/PiecewiseConstant_1/Greater (Greater) 
  SGDW/SGDW/update/PiecewiseConstant_1/Greater_1 (Greater) 
  SGDW/SGDW/update/PiecewiseConstant_1/LessEqual_1 (LessEqual) 
  SGDW/SGDW/update/PiecewiseConstant_1/and (LogicalAnd) 
  SGDW/SGDW/update/PiecewiseConstant_1/Greater_2 (Greater) 
  SGDW/SGDW/update/PiecewiseConstant_1/LessEqual_2 (LessEqual) 
  SGDW/SGDW/update/PiecewiseConstant_1/and_1 (LogicalAnd) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/preds_c (Pack) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/Cast (Cast) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/Const (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/num_true_conds (Sum) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/n_true_conds (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/LessEqual (LessEqual) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/Assert/Const (Const) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/Assert/AssertGuard/Identity (Identity) 
  SGDW/SGDW/update/PiecewiseConstant_1/case/cond/Identity (Identity) 
  SGDW/SGDW/update/Identity (Identity) 
  SGDW/SGDW/update/Identity_1 (Identity) 
  SGDW/SGDW/update/ResourceApplyKerasMomentum (ResourceApplyKerasMomentum) 

         [[{{node sequential/conv2d/Conv2D/ReadVariableOp}}]] [Op:__inference_train_function_3611]
```

/cc: @tanzhenyu 


Discussion:
/cc @MHStadler
@tanzhenyu @swghosh This does not seem to be an issue related to this recent fix - it is related to DecoupledWeightDecayExtension and PiecewiseConstantDecay
Other schedules work fine, but with PiecewiseConstantDecay it also fails if wd is kept constant, while lr is scheduled via PiecewiseConstantDecay. AdamW behaves equally, while normal tf.SGD works fine with this schedule.

See this colab: https://colab.research.google.com/drive/1ctd06tvN3O-4WEp-IXCsoEM2MR_odu4q?usp=sharing 
Okay, I managed to get a fix, looking at SGD (see FixedSGWD in the linked gist above)

The problem is that _resource_apply_dense and _resource_apply_sparse were missing the apply_state parameter - I will prepare a PR

I'm not yet sure why this only occurs for PiecewiseConstantDecay, I suspect it is related to how tf.case is executed on the GPU

Resolving patch:
diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers.py b/tensorflow_addons/optimizers/weight_decay_optimizers.py
--- a/tensorflow_addons/optimizers/weight_decay_optimizers.py
+++ b/tensorflow_addons/optimizers/weight_decay_optimizers.py
@@ -167,34 +167,59 @@ def apply_gradients(self, grads_and_vars, name=None, decay_var_list=None, **kwar
         )
         return super().apply_gradients(grads_and_vars, name=name, **kwargs)
 
-    def _decay_weights_op(self, var):
+    def _decay_weights_op(self, var, apply_state=None):
         if not self._decay_var_list or var.ref() in self._decay_var_list:
-            return var.assign_sub(self._decayed_wd(var.dtype) * var, self._use_locking)
+            var_device, var_dtype = var.device, var.dtype.base_dtype
+            coefficients = (apply_state or {}).get(
+                (var_device, var_dtype)
+            ) or self._fallback_apply_state(var_device, var_dtype)
+
+            return var.assign_sub(coefficients["wd_t"] * var, self._use_locking)
         return tf.no_op()
 
-    def _decay_weights_sparse_op(self, var, indices):
+    def _decay_weights_sparse_op(self, var, indices, apply_state=None):
         if not self._decay_var_list or var.ref() in self._decay_var_list:
-            update = -self._decayed_wd(var.dtype) * tf.gather(var, indices)
+            var_device, var_dtype = var.device, var.dtype.base_dtype
+            coefficients = (apply_state or {}).get(
+                (var_device, var_dtype)
+            ) or self._fallback_apply_state(var_device, var_dtype)
+
+            update = -coefficients["wd_t"] * tf.gather(var, indices)
             return self._resource_scatter_add(var, indices, update)
         return tf.no_op()
 
+    def _prepare_local(self, var_device, var_dtype, apply_state):
+        super(DecoupledWeightDecayExtension, self)._prepare_local(
+            var_device, var_dtype, apply_state
+        )
+
+        if "weight_decay" in self._hyper:
+            wd_t = tf.identity(self._decayed_wd(var_dtype))
+            apply_state[(var_device, var_dtype)]["wd_t"] = wd_t
+
     def _decayed_wd(self, var_dtype):
         wd_t = self._get_hyper("weight_decay", var_dtype)
+
         if isinstance(wd_t, tf.keras.optimizers.schedules.LearningRateSchedule):
             wd_t = tf.cast(wd_t(self.iterations), var_dtype)
+
         return wd_t
 
     # Here, we overwrite the apply functions that the base optimizer calls.
     # super().apply_x resolves to the apply_x function of the BaseOptimizer.
 
-    def _resource_apply_dense(self, grad, var):
-        with tf.control_dependencies([self._decay_weights_op(var)]):
-            return super()._resource_apply_dense(grad, var)
+    def _resource_apply_dense(self, grad, var, apply_state=None):
+        with tf.control_dependencies(
+            [self._decay_weights_op(var, apply_state=apply_state)]
+        ):
+            return super()._resource_apply_dense(grad, var, apply_state=apply_state)
 
-    def _resource_apply_sparse(self, grad, var, indices):
-        decay_op = self._decay_weights_sparse_op(var, indices)
+    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+        decay_op = self._decay_weights_sparse_op(var, indices, apply_state=apply_state)
         with tf.control_dependencies([decay_op]):
-            return super()._resource_apply_sparse(grad, var, indices)
+            return super()._resource_apply_sparse(
+                grad, var, indices, apply_state=apply_state
+            )
 
 
 @typechecked


Test patch:
diff --git a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
--- a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
+++ b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
@@ -266,6 +266,19 @@ def test_keras_fit_with_schedule():
     model.fit(x, y, epochs=1)
 
 
+@pytest.mark.with_device(["cpu", "gpu"])
+def test_weight_decay_with_piecewise_constant_decay_schedule():
+    model = tf.keras.models.Sequential([tf.keras.layers.Dense(2)])
+    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+    wd_schedule = tf.optimizers.schedules.PiecewiseConstantDecay([2], [1e-4, 1e-5])
+    optimizer = weight_decay_optimizers.SGDW(
+        learning_rate=1e-2, weight_decay=wd_schedule
+    )
+    model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])
+    x, y = np.random.uniform(size=(2, 4, 1))
+    model.fit(x, y, batch_size=1, epochs=1)
+
+
 @pytest.mark.parametrize("dtype", [(tf.half, 0), (tf.float32, 1), (tf.float64, 2)])
 def test_sparse_sgdw(dtype):
     do_test(