Repository: tensorflow/addons

Issue:
NovoGrad optimizer fails with embedding layer
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
- TensorFlow version and how it was installed (source or binary): TF 2.1 via pip
- TensorFlow-Addons version and how it was installed (source or binary): TFA 0.8.2 via pip
- Python version: 3.7.5
- Is GPU used? (yes/no): yes, CUDA 9.1

**Describe the bug**

Getting a fatal tensorflow error when NovoGrad is used. For some reason this doesn't occur when running on CPU machine, only when training on the GPU. Not sure where to start with debugging...

Traceback (most recent call last):
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 819, in fit
    use_multiprocessing=use_multiprocessing)
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 342, in fit
    total_epochs=epochs)
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 128, in run_one_epoch
    batch_outs = execution_function(iterator)
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 98, in execution_function
    distributed_function(input_fn))
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 568, in __call__
    result = self._call(*args, **kwds)
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 632, in _call
    return self._stateless_fn(*args, **kwds)
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 2363, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 1611, in _filtered_call
    self.captured_inputs)
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 1692, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 545, in call
    ctx=ctx)
  File "/home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  indices[3] = 4 is not in [0, 4)
	 [[node NovoGrad/NovoGrad/update_1/GatherV2 (defined at /home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_addons/optimizers/novograd.py:237) ]]
	 [[GroupCrossDeviceControlEdges_0/NovoGrad/NovoGrad/Const/_273]]
  (1) Invalid argument:  indices[3] = 4 is not in [0, 4)
	 [[node NovoGrad/NovoGrad/update_1/GatherV2 (defined at /home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_addons/optimizers/novograd.py:237) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_distributed_function_5894]

Errors may have originated from an input operation.

Input Source operations connected to node NovoGrad/NovoGrad/update_1/GatherV2:
 NovoGrad/NovoGrad/update_1/Unique (defined at CENSORED)	
 NovoGrad/NovoGrad/update_1/cond_2/Identity (defined at /home/jan/miniconda3/envs/tf/lib/python3.7/site-packages/tensorflow_addons/optimizers/novograd.py:230)

Function call stack:
distributed_function -> distributed_function

I use it for an RNN model built with the Model API, the model is compiled using sparse categorical crossentropy which is why I'm suspicious about the indexes in the error, this typically happens when an input is sent to the wrong embedding etc... Not sure where to start with debugging?



Discussion:
I suspect that tensorflow can't load CUDA or CUDNN but I'm unsure. Can you try with a simpler model? A dummy one with one layer? And does it work with a normal SGD? Do you get any warnings before the stacktrace?
Correction: my CUDA version is 10.1

I've tried with the simplest model: inputs - one LSTM layer - softmax, and the behavior is the same. Can confirm that this only happens with NovoGrad - I've tried Adam, RMSprop and LazyAdam from TFA without an issue.

There are some warnings before but they are the same as when using other optimizers.
```
2020-02-13 09:07:08.497207: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libnvinfer.so.6'; dlerror: libnvinfer.so.6: cannot open shared object file: No such file or directory
2020-02-13 09:07:08.497259: W tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libnvinfer_plugin.so.6'; dlerror: libnvinfer_plugin.so.6: cannot open shared object file: No such file or directory
2020-02-13 09:07:08.497266: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:30] Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2020-02-13 09:07:27.647490: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcuda.so.1
2020-02-13 09:07:27.652905: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.653187: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1555] Found device 0 with properties: 
pciBusID: 0000:2a:00.0 name: GeForce RTX 2080 Ti computeCapability: 7.5
coreClock: 1.665GHz coreCount: 68 deviceMemorySize: 10.75GiB deviceMemoryBandwidth: 573.69GiB/s
2020-02-13 09:07:27.653324: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudart.so.10.1
2020-02-13 09:07:27.656861: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcublas.so.10
2020-02-13 09:07:27.659169: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcufft.so.10
2020-02-13 09:07:27.659402: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcurand.so.10
2020-02-13 09:07:27.662103: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcusolver.so.10
2020-02-13 09:07:27.663642: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcusparse.so.10
2020-02-13 09:07:27.668007: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudnn.so.7
2020-02-13 09:07:27.668164: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.694166: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.694700: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1697] Adding visible gpu devices: 0
2020-02-13 09:07:27.695164: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2020-02-13 09:07:27.723545: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 3600115000 Hz
2020-02-13 09:07:27.724144: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x557ab39dd020 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-02-13 09:07:27.724178: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2020-02-13 09:07:27.857502: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.857957: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x557ab3603760 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2020-02-13 09:07:27.857990: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): GeForce RTX 2080 Ti, Compute Capability 7.5
2020-02-13 09:07:27.858256: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.858595: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1555] Found device 0 with properties: 
pciBusID: 0000:2a:00.0 name: GeForce RTX 2080 Ti computeCapability: 7.5
coreClock: 1.665GHz coreCount: 68 deviceMemorySize: 10.75GiB deviceMemoryBandwidth: 573.69GiB/s
2020-02-13 09:07:27.858644: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudart.so.10.1
2020-02-13 09:07:27.858656: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcublas.so.10
2020-02-13 09:07:27.858666: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcufft.so.10
2020-02-13 09:07:27.858677: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcurand.so.10
2020-02-13 09:07:27.858687: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcusolver.so.10
2020-02-13 09:07:27.858697: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcusparse.so.10
2020-02-13 09:07:27.858708: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudnn.so.7
2020-02-13 09:07:27.858765: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.859088: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.859378: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1697] Adding visible gpu devices: 0
2020-02-13 09:07:27.859421: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudart.so.10.1
2020-02-13 09:07:27.860205: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1096] Device interconnect StreamExecutor with strength 1 edge matrix:
2020-02-13 09:07:27.860215: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1102]      0 
2020-02-13 09:07:27.860222: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] 0:   N 
2020-02-13 09:07:27.860338: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.860676: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:981] 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-02-13 09:07:27.860962: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:39] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.
2020-02-13 09:07:27.861005: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1241] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 2742 MB memory) -> physical GPU (device: 0, name: GeForce RTX 2080 Ti, pci bus id: 0000:2a:00.0, compute capability: 7.5)
./logs/durables_demo Thu Feb 13 09:07:29 2020
2020-02-13 09:07:44.683722: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcublas.so.10
2020-02-13 09:07:45.541001: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudnn.so.7
2020-02-13 09:07:47.157708: W tensorflow/core/common_runtime/base_collective_executor.cc:217] BaseCollectiveExecutor::StartAbort Invalid argument: indices[4] = 36 is not in [0, 19)
	 [[{{node Lookahead/Lookahead/update_1/GatherV2}}]]
	 [[Lookahead/Lookahead/update_17/cond_1/then/_521/mul/ReadVariableOp/_1888]]
2020-02-13 09:07:47.161827: W tensorflow/core/common_runtime/base_collective_executor.cc:217] BaseCollectiveExecutor::StartAbort Invalid argument: indices[4] = 36 is not in [0, 19)
	 [[{{node Lookahead/Lookahead/update_1/GatherV2}}]]
2020-02-13 09:07:47.184974: I tensorflow/core/profiler/lib/profiler_session.cc:225] Profiler session started.
2020-02-13 09:07:47.185182: I tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1259] Profiler found 1 GPUs
2020-02-13 09:07:47.189878: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcupti.so.10.1
2020-02-13 09:07:47.291371: E tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1307] function cupti_interface_->Subscribe( &subscriber_, (CUpti_CallbackFunc)ApiCallback, this)failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
2020-02-13 09:07:47.292672: E tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1346] function cupti_interface_->ActivityRegisterCallbacks( AllocCuptiActivityBuffer, FreeCuptiActivityBuffer)failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
2020-02-13 09:07:47.307569: E tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1329] function cupti_interface_->EnableCallback( 0 , subscriber_, CUPTI_CB_DOMAIN_DRIVER_API, cbid)failed with error CUPTI_ERROR_INVALID_PARAMETER
2020-02-13 09:07:47.307604: I tensorflow/core/profiler/internal/gpu/device_tracer.cc:88]  GpuTracer has collected 0 callback api events and 0 activity events.
```

UPDATE: I have found the culprit: I'm using Embedding layers for all the inputs and somehow NovoGrad gets confused. Without Embeddings the error doesn't occur. Does that help?
Thanks! That helps a lot! @shreyashpatodia could you take it from here? Thank you very much! 
Hi @valendin, could you provide the code snippet of your model and training script? Thank you.
Here's an example that breaks - if you replace NovoGrad with Adam, everything works. 

```
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
import tensorflow_addons as tfa
import tensorflow.keras.backend as K
from tensorflow.keras import Model
from tensorflow.keras.layers import Input, Embedding, LSTM, Dense, Concatenate, Lambda

seq_len = 3
max_input = 10

input_one = Input(shape=(seq_len, 1), name='one')
input_two = Input(shape=(seq_len, 1), name='two')

embedding_one = Embedding(input_dim=max_input+1, output_dim=5)
embedding_two = Embedding(input_dim=max_input+1, output_dim=5)

concat = Concatenate(axis=-1)
squeeze = Lambda(lambda x: K.squeeze(x, axis=-2))

lstm_layer = LSTM(32, return_sequences=True, stateful=False)
dense_layer = Dense(32)
softmax_layer = Dense(2, activation='softmax')

emb_one = embedding_one(input_one)
emb_two = embedding_two(input_two)
emb_one = squeeze(emb_one)
emb_two = squeeze(emb_two)
output = concat([emb_one, emb_two])
output = lstm_layer(output)
output = dense_layer(output)
output = softmax_layer(output)

inputs = [input_one, input_two]
model = Model(inputs, output)

from tensorflow.keras.losses import sparse_categorical_crossentropy
from tensorflow.keras.optimizers import Adam
from tensorflow_addons.optimizers import NovoGrad

optimizer = NovoGrad()

model.compile(loss=sparse_categorical_crossentropy, optimizer=optimizer)

def decode_sample(sample, target):
    return ({'one'      : tf.cast(tf.expand_dims(sample[:, 0], axis=-1), 'int32'), 
             'two'      : tf.cast(tf.expand_dims(sample[:, 1], axis=-1), 'int32')},
             tf.cast(tf.expand_dims(target, axis=-1), 'int32'))

train_x = [np.array([[0, max_input],[0, 0], [0, max_input]]) for _ in range(10)]
train_y = [np.array([[0],[1],[0]]) for _ in range(10)]
dataset = (tf.data.Dataset.from_tensor_slices((train_x, train_y))
                .map(decode_sample)
                .batch(2)
                .repeat())

model.summary()

history = model.fit(dataset, 
                    epochs=2, 
                    verbose=1,
                    shuffle=True,
                    steps_per_epoch=5)
```

Resolving patch:
diff --git a/tensorflow_addons/optimizers/novograd.py b/tensorflow_addons/optimizers/novograd.py
--- a/tensorflow_addons/optimizers/novograd.py
+++ b/tensorflow_addons/optimizers/novograd.py
@@ -222,7 +222,9 @@ def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
         else:
             grad = grad / (tf.sqrt(v_t) + self.epsilon)
         grad = tf.cond(
-            tf.greater(weight_decay, 0), lambda: grad + weight_decay * var, lambda: grad
+            tf.greater(weight_decay, 0),
+            lambda: grad + weight_decay * tf.gather(var, indices),
+            lambda: grad,
         )
         grad = tf.cond(
             tf.logical_and(grad_averaging, tf.not_equal(self.iterations, 0)),
@@ -234,7 +236,7 @@ def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
             var.handle,
             m.handle,
             coefficients["lr_t"],
-            tf.gather(grad, indices),
+            grad,
             indices,
             coefficients["beta_1_t"],
             use_locking=self._use_locking,


Test patch:
diff --git a/tensorflow_addons/optimizers/tests/novograd_test.py b/tensorflow_addons/optimizers/tests/novograd_test.py
--- a/tensorflow_addons/optimizers/tests/novograd_test.py
+++ b/tensorflow_addons/optimizers/tests/novograd_test.py
@@ -68,27 +68,23 @@ def run_sparse_sample(iterations, expected, optimizer):
     var_0 = tf.Variable([1.0, 2.0])
     var_1 = tf.Variable([3.0, 4.0])
 
-    grad_0 = tf.IndexedSlices(
-        tf.constant([0.1, 0.2]), tf.constant([0, 1]), tf.constant([2])
-    )
-    grad_1 = tf.IndexedSlices(
-        tf.constant([0.3, 0.4]), tf.constant([0, 1]), tf.constant([2])
-    )
+    grad_0 = tf.IndexedSlices(tf.constant([0.1]), tf.constant([0]), tf.constant([2]))
+    grad_1 = tf.IndexedSlices(tf.constant([0.4]), tf.constant([1]), tf.constant([2]))
 
     grads_and_vars = list(zip([grad_0, grad_1], [var_0, var_1]))
 
     for _ in range(iterations):
         optimizer.apply_gradients(grads_and_vars)
 
-    np.testing.assert_allclose(var_0.read_value(), expected[0], atol=2e-4)
-    np.testing.assert_allclose(var_1.read_value(), expected[1], atol=2e-4)
+    np.testing.assert_allclose(var_0.read_value(), expected[0])
+    np.testing.assert_allclose(var_1.read_value(), expected[1])
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_sparse_sample():
     run_sparse_sample(
-        iterations=1,
-        expected=[[0.9552786425, 1.9105572849], [2.9400000012, 3.9200000016]],
+        iterations=2,
+        expected=[[0.71, 2.0], [3.0, 3.71]],
         optimizer=NovoGrad(lr=0.1, epsilon=1e-8),
     )
 
@@ -96,8 +92,8 @@ def test_sparse_sample():
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_sparse_sample_with_weight_decay():
     run_sparse_sample(
-        iterations=1,
-        expected=[[0.945278642, 1.8905572849], [2.9100000012, 3.8800000016]],
+        iterations=2,
+        expected=[[0.6821, 2.0], [3.0, 3.5954]],
         optimizer=NovoGrad(lr=0.1, weight_decay=0.1, epsilon=1e-8),
     )
 
@@ -106,7 +102,7 @@ def test_sparse_sample_with_weight_decay():
 def test_sparse_sample_with_grad_averaging():
     run_sparse_sample(
         iterations=2,
-        expected=[[0.9105572849, 1.8211145698], [2.8800000024, 3.8400000032]],
+        expected=[[0.8, 2.0], [3.0, 3.8]],
         optimizer=NovoGrad(lr=0.1, grad_averaging=True, epsilon=1e-8),
     )