Add tracing API to the Python package.

PiperOrigin-RevId: 342128507

tensorflow_federated/BUILD

@@ -26,6 +26,7 @@ py_library(
     deps = [
         ":version",
         "//tensorflow_federated/python/aggregators",
+        "//tensorflow_federated/python/common_libs:tracing",
         "//tensorflow_federated/python/core/api:computation_base",
         "//tensorflow_federated/python/core/api:computation_types",
         "//tensorflow_federated/python/core/api:computations",

tensorflow_federated/__init__.py

@@ -20,6 +20,7 @@ from tensorflow_federated.version import __version__  # pylint: disable=g-bad-im
 from tensorflow_federated.python import aggregators
 from tensorflow_federated.python import learning
 from tensorflow_federated.python import simulation
+from tensorflow_federated.python.common_libs import tracing as profiler
 from tensorflow_federated.python.core import backends
 from tensorflow_federated.python.core import framework
 from tensorflow_federated.python.core import templates

tensorflow_federated/python/common_libs/BUILD

@@ -126,8 +126,10 @@ py_library(
     name = "tracing",
     srcs = ["tracing.py"],
     srcs_version = "PY3",
-    # TODO(b/163856734): Fix the visibility of the `common_libs`.
-    visibility = ["//tensorflow_federated:internal"],
+    visibility = [
+        "//tensorflow_federated:__pkg__",
+        "//tensorflow_federated:internal",
+    ],
     deps = [":py_typecheck"],
 )