Repository: tensorflow/addons

Issue:
NovoGrad optimizer doesn't work with float64 backend setting
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Debian 10 Buster
- TensorFlow version and how it was installed (source or binary): 2.4.1 (pip install tensorflow)
- TensorFlow-Addons version and how it was installed (source or binary): 0.12.1 (pip install)
- Python version: 3.8.9
- Is GPU used? (yes/no): no

**Describe the bug**

NovoGrad optimizer doesn't work when backend is set to use float64 by default.

**Code to reproduce the issue**

Just run [NovoGrad test](https://github.com/tensorflow/addons/blob/v0.12.1/tensorflow_addons/optimizers/tests/novograd_test.py#L110) with different settings:

```
import numpy as np
import tensorflow as tf
from tensorflow_addons.optimizers import NovoGrad


def test_fit_simple_linear_model():
    np.random.seed(0x2020)
    tf.random.set_seed(0x2020)

    x = np.random.standard_normal((100000, 3))
    w = np.random.standard_normal((3, 1))
    y = np.dot(x, w) + np.random.standard_normal((100000, 1)) * 1e-5

    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Dense(input_shape=(3,), units=1))
    model.compile(NovoGrad(), loss="mse")

    model.fit(x, y, epochs=2)

    x = np.random.standard_normal((100, 3))
    y = np.dot(x, w)
    predicted = model.predict(x)

    max_abs_diff = np.max(np.abs(predicted - y))
    assert max_abs_diff < 1e-2
    
tf.keras.backend.set_floatx('float32')
test_fit_simple_linear_model()
tf.keras.backend.set_floatx('float64')
test_fit_simple_linear_model()
```

**Other info / logs**

Output:
```
Epoch 1/2
3125/3125 [==============================] - 3s 705us/step - loss: 1.0213
Epoch 2/2
3125/3125 [==============================] - 2s 712us/step - loss: 1.0537e-10
Epoch 1/2
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-c5171e788b94> in <module>
     28 test_fit_simple_linear_model()
     29 tf.keras.backend.set_floatx('float64')
---> 30 test_fit_simple_linear_model()

<ipython-input-5-c5171e788b94> in test_fit_simple_linear_model()
     16     model.compile(NovoGrad(), loss="mse")
     17 
---> 18     model.fit(x, y, epochs=2)
     19 
     20     x = np.random.standard_normal((100, 3))

/usr/local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
   1098                 _r=1):
   1099               callbacks.on_train_batch_begin(step)
-> 1100               tmp_logs = self.train_function(iterator)
   1101               if data_handler.should_sync:
   1102                 context.async_wait()

/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    826     tracing_count = self.experimental_get_tracing_count()
    827     with trace.Trace(self._name) as tm:
--> 828       result = self._call(*args, **kwds)
    829       compiler = "xla" if self._experimental_compile else "nonXla"
    830       new_tracing_count = self.experimental_get_tracing_count()

/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    869       # This is the first call of __call__, so we have to initialize.
    870       initializers = []
--> 871       self._initialize(args, kwds, add_initializers_to=initializers)
    872     finally:
    873       # At this point we know that the initialization is complete (or less

/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    723     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
    724     self._concrete_stateful_fn = (
--> 725         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
    726             *args, **kwds))
    727 

/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   2967       args, kwargs = None, None
   2968     with self._lock:
-> 2969       graph_function, _ = self._maybe_define_function(args, kwargs)
   2970     return graph_function
   2971 

/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3359 
   3360           self._function_cache.missed.add(call_context_key)
-> 3361           graph_function = self._create_graph_function(args, kwargs)
   3362           self._function_cache.primary[cache_key] = graph_function
   3363 

/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3194     arg_names = base_arg_names + missing_arg_names
   3195     graph_function = ConcreteFunction(
-> 3196         func_graph_module.func_graph_from_py_func(
   3197             self._name,
   3198             self._python_function,

/usr/local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    988         _, original_func = tf_decorator.unwrap(python_func)
    989 
--> 990       func_outputs = python_func(*func_args, **func_kwargs)
    991 
    992       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    632             xla_context.Exit()
    633         else:
--> 634           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    635         return out
    636 

/usr/local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    975           except Exception as e:  # pylint:disable=broad-except
    976             if hasattr(e, "ag_error_metadata"):
--> 977               raise e.ag_error_metadata.to_exception(e)
    978             else:
    979               raise

TypeError: in user code:

    /usr/local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:805 train_function  *
        return step_function(self, iterator)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:795 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:1259 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2730 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:3417 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:788 run_step  **
        outputs = model.train_step(data)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:757 train_step
        self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:498 minimize
        return self.apply_gradients(grads_and_vars, name=name)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:631 apply_gradients
        return distribute_ctx.get_replica_context().merge_call(
    /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2941 merge_call
        return self._merge_call(merge_fn, args, kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2948 _merge_call
        return merge_fn(self._strategy, *args, **kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:682 _distributed_apply  **
        update_ops.extend(distribution.extended.update(
    /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2494 update
        return self._update(var, fn, args, kwargs, group)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:3431 _update
        return self._update_non_slot(var, fn, (var,) + tuple(args), kwargs, group)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:3437 _update_non_slot
        result = fn(*args, **kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:658 apply_grad_to_update_var  **
        update_op = self._resource_apply_dense(grad, var, **apply_kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow_addons/optimizers/novograd.py:163 _resource_apply_dense
        v_t = tf.cond(
    /usr/local/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/control_flow_ops.py:1396 cond_for_tf_v2
        return cond(pred, true_fn=true_fn, false_fn=false_fn, strict=True, name=name)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/util/deprecation.py:538 new_func
        return func(*args, **kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/control_flow_ops.py:1180 cond
        return cond_v2.cond_v2(pred, true_fn, false_fn, name)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/cond_v2.py:90 cond_v2
        false_graph = func_graph_module.func_graph_from_py_func(
    /usr/local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py:990 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow_addons/optimizers/novograd.py:167 <lambda>
        + g_2 * coefficients["one_minus_beta_2_t"],
    /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py:1180 binary_op_wrapper
        raise e
    /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py:1164 binary_op_wrapper
        return func(x, y, name=name)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py:1496 _mul_dispatch
        return multiply(x, y, name=name)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py:518 multiply
        return gen_math_ops.mul(x, y, name)
    /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py:6077 mul
        _, _, _op, _outputs = _op_def_library._apply_op_helper(
    /usr/local/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py:555 _apply_op_helper
        raise TypeError(

    TypeError: Input 'y' of 'Mul' Op has type float64 that does not match type float32 of argument 'x'.
```

I suspect that's because NovoGrad casts gradients to float32 - https://github.com/tensorflow/addons/blob/v0.12.1/tensorflow_addons/optimizers/novograd.py#L162
It can be better to replace `tf.float32` with `tf.keras.backend.floatx()` here and in other similar places.


Discussion:
Can you prepare a PR?

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
@@ -155,11 +155,11 @@ def _resource_apply_dense(self, grad, var, apply_state=None):
         coefficients = (apply_state or {}).get(
             (var_device, var_dtype)
         ) or self._fallback_apply_state(var_device, var_dtype)
-        weight_decay = self._get_hyper("weight_decay")
+        weight_decay = self._get_hyper("weight_decay", var_dtype)
         grad_averaging = self._get_hyper("grad_averaging")
 
         v = self.get_slot(var, "v")
-        g_2 = tf.reduce_sum(tf.square(tf.cast(grad, tf.float32)))
+        g_2 = tf.reduce_sum(tf.square(grad))
         v_t = tf.cond(
             tf.equal(self.iterations, 0),
             lambda: g_2,
@@ -198,11 +198,11 @@ def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
         coefficients = (apply_state or {}).get(
             (var_device, var_dtype)
         ) or self._fallback_apply_state(var_device, var_dtype)
-        weight_decay = self._get_hyper("weight_decay")
+        weight_decay = self._get_hyper("weight_decay", var_dtype)
         grad_averaging = self._get_hyper("grad_averaging")
 
         v = self.get_slot(var, "v")
-        g_2 = tf.reduce_sum(tf.square(tf.cast(grad, tf.float32)))
+        g_2 = tf.reduce_sum(tf.square(grad))
         # v is just a scalar and does not need to involve sparse tensors.
         v_t = tf.cond(
             tf.equal(self.iterations, 0),


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
@@ -21,12 +21,12 @@
 from tensorflow_addons.optimizers import NovoGrad
 
 
-def run_dense_sample(iterations, expected, optimizer):
-    var_0 = tf.Variable([1.0, 2.0], dtype=tf.dtypes.float32)
-    var_1 = tf.Variable([3.0, 4.0], dtype=tf.dtypes.float32)
+def run_dense_sample(iterations, expected, optimizer, dtype):
+    var_0 = tf.Variable([1.0, 2.0], dtype=dtype)
+    var_1 = tf.Variable([3.0, 4.0], dtype=dtype)
 
-    grad_0 = tf.constant([0.1, 0.2], dtype=tf.dtypes.float32)
-    grad_1 = tf.constant([0.3, 0.4], dtype=tf.dtypes.float32)
+    grad_0 = tf.constant([0.1, 0.2], dtype=dtype)
+    grad_1 = tf.constant([0.3, 0.4], dtype=dtype)
 
     grads_and_vars = list(zip([grad_0, grad_1], [var_0, var_1]))
 
@@ -38,38 +38,48 @@ def run_dense_sample(iterations, expected, optimizer):
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-def test_dense_sample():
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_dense_sample(dtype):
     run_dense_sample(
         iterations=1,
         expected=[[0.9552786425, 1.9105572849], [2.9400000012, 3.9200000016]],
         optimizer=NovoGrad(lr=0.1, epsilon=1e-8),
+        dtype=dtype,
     )
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-def test_dense_sample_with_weight_decay():
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_dense_sample_with_weight_decay(dtype):
     run_dense_sample(
         iterations=1,
         expected=[[0.945278642, 1.8905572849], [2.9100000012, 3.8800000016]],
         optimizer=NovoGrad(lr=0.1, weight_decay=0.1, epsilon=1e-8),
+        dtype=dtype,
     )
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-def test_dense_sample_with_grad_averaging():
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_dense_sample_with_grad_averaging(dtype):
     run_dense_sample(
         iterations=2,
         expected=[[0.9105572849, 1.8211145698], [2.8800000024, 3.8400000032]],
         optimizer=NovoGrad(lr=0.1, grad_averaging=True, epsilon=1e-8),
+        dtype=dtype,
     )
 
 
-def run_sparse_sample(iterations, expected, optimizer):
-    var_0 = tf.Variable([1.0, 2.0])
-    var_1 = tf.Variable([3.0, 4.0])
+def run_sparse_sample(iterations, expected, optimizer, dtype):
+    var_0 = tf.Variable([1.0, 2.0], dtype=dtype)
+    var_1 = tf.Variable([3.0, 4.0], dtype=dtype)
 
-    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]))
+    grad_0 = tf.IndexedSlices(
+        tf.constant([0.1], dtype=dtype), tf.constant([0]), tf.constant([2])
+    )
+    grad_1 = tf.IndexedSlices(
+        tf.constant([0.4], dtype=dtype), tf.constant([1]), tf.constant([2])
+    )
 
     grads_and_vars = list(zip([grad_0, grad_1], [var_0, var_1]))
 
@@ -81,29 +91,35 @@ def run_sparse_sample(iterations, expected, optimizer):
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-def test_sparse_sample():
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparse_sample(dtype):
     run_sparse_sample(
         iterations=2,
         expected=[[0.71, 2.0], [3.0, 3.71]],
         optimizer=NovoGrad(lr=0.1, epsilon=1e-8),
+        dtype=dtype,
     )
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-def test_sparse_sample_with_weight_decay():
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparse_sample_with_weight_decay(dtype):
     run_sparse_sample(
         iterations=2,
         expected=[[0.6821, 2.0], [3.0, 3.5954]],
         optimizer=NovoGrad(lr=0.1, weight_decay=0.1, epsilon=1e-8),
+        dtype=dtype,
     )
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-def test_sparse_sample_with_grad_averaging():
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparse_sample_with_grad_averaging(dtype):
     run_sparse_sample(
         iterations=2,
         expected=[[0.8, 2.0], [3.0, 3.8]],
         optimizer=NovoGrad(lr=0.1, grad_averaging=True, epsilon=1e-8),
+        dtype=dtype,
     )