# Buggy code
def apply_gradients(self, grads_and_vars, name=None):
    self._optimizer._iterations = self.iterations  # pylint: disable=protected-access
    return super().apply_gradients(grads_and_vars, name)


# Traceback (TypeError)
Epoch 1/5
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-19-841302a83165> in <module>
----> 1 model.fit(
2     train_ds,
3     epochs=1 if lr_finder else 5,
4     callbacks=callbacks,
5     steps_per_epoch=findlr_steps if lr_finder else None,

~/.local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py in _method_wrapper(self, *args, **kwargs)
64   def _method_wrapper(self, *args, **kwargs):
65     if not self._in_multi_worker_mode():  # pylint: disable=protected-access
---> 66       return method(self, *args, **kwargs)
67
68     # Running inside `run_distribute_coordinator` already.

~/.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)
846                 batch_size=batch_size):
847               callbacks.on_train_batch_begin(step)
--> 848               tmp_logs = train_function(iterator)
849               # Catch OutOfRangeError for Datasets of unknown size.
850               # This blocks until the batch has finished executing.

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
578         xla_context.Exit()
579     else:
--> 580       result = self._call(*args, **kwds)
581
582     if tracing_count == self._get_tracing_count():

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
625       # This is the first call of __call__, so we have to initialize.
626       initializers = []
--> 627       self._initialize(args, kwds, add_initializers_to=initializers)
628     finally:
629       # At this point we know that the initialization is complete (or less

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
503     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
504     self._concrete_stateful_fn = (
--> 505         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
506             *args, **kwds))
507

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
2444       args, kwargs = None, None
2445     with self._lock:
-> 2446       graph_function, _, _ = self._maybe_define_function(args, kwargs)
2447     return graph_function
2448

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
2775
2776       self._function_cache.missed.add(call_context_key)
-> 2777       graph_function = self._create_graph_function(args, kwargs)
2778       self._function_cache.primary[cache_key] = graph_function
2779       return graph_function, args, kwargs

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
2655     arg_names = base_arg_names + missing_arg_names
2656     graph_function = ConcreteFunction(
-> 2657         func_graph_module.func_graph_from_py_func(
2658             self._name,
2659             self._python_function,

~/.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)
979         _, original_func = tf_decorator.unwrap(python_func)
980
--> 981       func_outputs = python_func(*func_args, **func_kwargs)
982
983       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
439         # __wrapped__ allows AutoGraph to swap in a converted function. We give
440         # the function a weak reference to itself to avoid a reference cycle.
--> 441         return weak_wrapped_fn().__wrapped__(*args, **kwds)
442     weak_wrapped_fn = weakref.ref(wrapped_fn)
443

~/.local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
966           except Exception as e:  # pylint:disable=broad-except
967             if hasattr(e, "ag_error_metadata"):
--> 968               raise e.ag_error_metadata.to_exception(e)
969             else:
970               raise

TypeError: in user code:

/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:571 train_function  *
outputs = self.distribute_strategy.run(
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:951 run  **
return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2290 call_for_each_replica
return self._call_for_each_replica(fn, args, kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2649 _call_for_each_replica
return fn(*args, **kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:540 train_step  **
_minimize(self.distribute_strategy, tape, self.optimizer, loss,
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1810 _minimize
optimizer.apply_gradients(
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer.py:245 apply_gradients
return distribution_strategy_context.get_replica_context().merge_call(
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2420 merge_call
return self._merge_call(merge_fn, args, kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2427 _merge_call
return merge_fn(self._strategy, *args, **kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer.py:269 _apply_gradients_cross_replica  **
maybe_apply_op = smart_cond.smart_cond(should_apply_grads,
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py:58 smart_cond
return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/util/deprecation.py:507 new_func
return func(*args, **kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/ops/control_flow_ops.py:1177 cond
return cond_v2.cond_v2(pred, true_fn, false_fn, name)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/ops/cond_v2.py:78 cond_v2
true_graph = func_graph_module.func_graph_from_py_func(
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py:981 func_graph_from_py_func
func_outputs = python_func(*func_args, **func_kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer.py:261 apply_fn
return distribution.extended.call_for_each_replica(
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2290 call_for_each_replica
return self._call_for_each_replica(fn, args, kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2649 _call_for_each_replica
return fn(*args, **kwargs)
/home/alex/.local/lib/python3.8/site-packages/tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer.py:279 _apply_gradients
return self._optimizer.apply_gradients(

TypeError: apply_gradients() got an unexpected keyword argument 'experimental_aggregate_gradients'

# Fixed code
def apply_gradients(self, grads_and_vars, name=None, **kwargs):
    self._optimizer._iterations = self.iterations  # pylint: disable=protected-access
    return super().apply_gradients(grads_and_vars, name, **kwargs)