Add native_module.get_unsupported_collections().

PiperOrigin-RevId: 233584508

tensorflow_hub/native_module.py

@@ -709,6 +709,10 @@ def check_collections_are_supported(saved_model_handler, supported):
                        " as appropriate." % list(unsupported))
 
 
+def get_unsupported_collections(used_collection_keys):
+  return list(set(used_collection_keys) - _SUPPORTED_COLLECTIONS)
+
+
 def register_ops_if_needed(graph_ops):
   """Register graph ops absent in op_def_registry, if present in c++ registry.
 

tensorflow_hub/native_module_test.py

@@ -129,6 +129,26 @@ def wrong_module_fn():
       spec = native_module.create_module_spec(wrong_module_fn)
     self.assertIn("No signatures present", str(cm.exception))
 
+  def testUnsupportedCollections(self):
+
+    def module_fn():
+      scale = tf.get_variable("x", (), collections=["my_scope"])
+      x = tf.placeholder(tf.float32, shape=[None, 3])
+      native_module.add_signature("my_func", {"x": x}, {"y": x*scale})
+
+    with self.assertRaises(ValueError) as cm:
+      _ = native_module.create_module_spec(module_fn)
+      self.assertIn("Unsupported collections in graph", cm)
+
+    with tf.Graph().as_default() as tmp_graph:
+      module_fn()
+      unsupported_collections = native_module.get_unsupported_collections(
+          tmp_graph.get_all_collection_keys())
+      self.assertEqual(["my_scope"], unsupported_collections)
+
+    _ = native_module.create_module_spec(
+        module_fn, drop_collections=unsupported_collections)
+
 
 class RecoverPartitionedVariableMapTest(tf.test.TestCase):