[fix] Support TF CheckpointManager in 2.9.

User may create Checkpoint by passing keys "model", "optimizer" and so on. So root parameter in __init__ function may be None. We need to walk through all kwargs.

Author: MoFHeka

tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py

@@ -339,7 +339,7 @@ def common_all_to_all_embedding_trainable_v2(self, base_opt, test_opt, name):
             bp_v2=False,
             kv_creator=kv_creator,
             name='all2all_emb')
-        ckpt = de.train.DEHvdCheckpoint(new_base_model)
+        ckpt = de.train.DEHvdCheckpoint(my_model=new_base_model)
         hvd.join()  # Sync for avoiding files conflict
         ckpt.restore(tf.train.latest_checkpoint(save_dir + '/ckpt/'))
         new_a2aemb_size = new_base_model.layers[0].params.size()

tensorflow_recommenders_addons/dynamic_embedding/python/train/checkpoint.py

@@ -155,17 +155,21 @@ def _filter_de_hvd_a2a_tw(var):
         return False
       return True
 
-    if _filter_de_hvd_a2a_tw(self.root):
-      func(var.params, de_dir)
-    if hasattr(self.root, 'variables'):
-      for var in self.root.variables:
-        if _filter_de_hvd_a2a_tw(var):
-          func(var.params, de_dir)
+    def _handle_model_or_variable(obj):
+      if _filter_de_hvd_a2a_tw(obj):
+        func(var.params, de_dir)
+      if hasattr(obj, 'variables'):
+        _iter = obj.variables() if callable(obj.variables) else obj.variables
+        for var in _iter:
+          if _filter_de_hvd_a2a_tw(var):
+            func(var.params, de_dir)
+
+    if hasattr(self, 'root'):
+      _handle_model_or_variable(self.root)
     if len(self._tmp_var_key_set):
-      for var_key in self._tmp_var_key_set:
-        var = getattr(self, var_key)
-        if _filter_de_hvd_a2a_tw(var):
-          func(var.params, de_dir)
+      for obj_key in self._tmp_var_key_set:
+        obj_var = getattr(self, obj_key)
+        _handle_model_or_variable(obj_var)
 
   def _de_hvd_write_fs_func(self, file_prefix, tf_write_func):