In models/README.md, recommend MtAlbis.

PiperOrigin-RevId: 617512268
tensorflow · Mar 20, 2024 · 84c5400 · 84c5400
1 parent 90331a1
commit 84c5400
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 16 deletions.
diff --git a/tensorflow_gnn/models/README.md b/tensorflow_gnn/models/README.md
@@ -44,3 +44,6 @@ versioning](https://semver.org/spec/v2.0.0.html) of the TF-GNN package.
     multi-head attention on graph (Dwivedi&Bresson, 2021).
   * [VanillaMPNN](vanilla_mpnn/README.md): TF-GNN's classic baseline model,
     based on (Gilmer&al., 2016).
+
+Unsure? For generic node prediction tasks on relational data, we recommend
+to start with MtAlbis.
diff --git a/tensorflow_gnn/models/gcn/BUILD b/tensorflow_gnn/models/gcn/BUILD
@@ -7,12 +7,7 @@ package(
     default_visibility = [":__subpackages__"],
 )
 
-package_group(
-    name = "users",
-    packages = [
-        "//third_party/google_research/google_research/ugsl/...",
-    ],
-)
+package_group(name = "users")
 
 pytype_strict_library(
     name = "gcn",

diff --git a/tensorflow_gnn/models/mt_albis/BUILD b/tensorflow_gnn/models/mt_albis/BUILD
@@ -7,16 +7,12 @@ package(
     default_visibility = [":__subpackages__"],
 )
 
-package_group(name = "users")
-
+# This is the public BUILD target for users of this model.
 pytype_strict_library(
     name = "mt_albis",
     srcs = ["__init__.py"],
     srcs_version = "PY3",
-    visibility = [
-        ":__subpackages__",
-        ":users",
-    ],
+    visibility = ["//visibility:public"],
     deps = [
         ":config_dict",
         ":layers",
@@ -29,10 +25,7 @@ pytype_strict_library(
 pytype_strict_library(
     name = "hparams_vizier",
     srcs = ["hparams_vizier.py"],
-    visibility = [
-        ":__subpackages__",
-        ":users",
-    ],
+    visibility = ["//visibility:public"],
     deps = [
         "//:expect_vizier_service_pyvizier_installed",
     ],