Remove tfp.experimental.substrates. Users should use

tfp.substrates.{jax,numpy} instead.

PiperOrigin-RevId: 628430022

Author: ThomasColthurst

BUILD

@@ -20,5 +20,6 @@ sh_binary(
         "required_packages.py",
         "setup.py",
         "//tensorflow_probability",
+        "//tensorflow_probability/substrates",
     ],
 )

conftest.py

@@ -20,7 +20,6 @@
 collect_ignore = [
     "discussion/",
     "setup.py",
-    "tensorflow_probability/python/experimental/substrates/"
 ]
 
 

tensorflow_probability/python/BUILD

@@ -41,6 +41,7 @@ py_library(
         "//tensorflow_probability/python/glm",
         "//tensorflow_probability/python/internal",
         "//tensorflow_probability/python/internal:all_util",
+        "//tensorflow_probability/python/internal:lazy_loader",
         "//tensorflow_probability/python/layers",
         "//tensorflow_probability/python/math",
         "//tensorflow_probability/python/mcmc",

tensorflow_probability/python/experimental/BUILD

@@ -41,7 +41,6 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/experimental/auto_batching",
         "//tensorflow_probability/python/experimental/marginalize",
         "//tensorflow_probability/python/experimental/nn",
-        "//tensorflow_probability/python/experimental/substrates",
         "//tensorflow_probability/python/experimental/timeseries",
         "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/experimental/util:composite_tensor",
@@ -64,7 +63,6 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/experimental/sequential",
         "//tensorflow_probability/python/experimental/stats",
         "//tensorflow_probability/python/experimental/sts_gibbs",
-        "//tensorflow_probability/python/experimental/substrates",
         "//tensorflow_probability/python/experimental/tangent_spaces",
         "//tensorflow_probability/python/experimental/timeseries",
         "//tensorflow_probability/python/experimental/util",

tensorflow_probability/python/experimental/__init__.py

@@ -43,7 +43,6 @@
 from tensorflow_probability.python.experimental import sequential
 from tensorflow_probability.python.experimental import stats
 from tensorflow_probability.python.experimental import sts_gibbs
-from tensorflow_probability.python.experimental import substrates
 from tensorflow_probability.python.experimental import tangent_spaces
 from tensorflow_probability.python.experimental import timeseries
 from tensorflow_probability.python.experimental import util
@@ -76,7 +75,6 @@
     'sequential',
     'sts_gibbs',
     'stats',
-    'substrates',
     'tangent_spaces',
     'timeseries',
     'unnest',

tensorflow_probability/python/experimental/substrates/BUILD

@@ -19,6 +19,9 @@
 import types
 
 
+__all__ = ['LazyLoader']
+
+
 class LazyLoader(types.ModuleType):
   """Lazily import a module to avoid pulling in large deps, defer checks."""
 

tensorflow_probability/python/experimental/substrates/__init__.py

@@ -32,6 +32,8 @@ py_library(
     srcs = [
         "__init__.py",
     ],
+    # :substrates needs to be visible to the external :pip_pkg target.
+    visibility = ["//visibility:public"],  # EnableOnExport
     deps = [
         ":jax",
         ":numpy",

tensorflow_probability/python/internal/lazy_loader.py

@@ -258,6 +258,9 @@ def disable_all(name):
   if FLAGS.numpy_to_jax:
     contents = contents.replace('tfp.substrates.numpy', 'tfp.substrates.jax')
     contents = contents.replace('substrates.numpy', 'substrates.jax')
+    contents = contents.replace(
+        'tensorflow_probability.substrates import numpy',
+        'tensorflow_probability.substrates import jax')
     contents = contents.replace('backend.numpy', 'backend.jax')
     contents = contents.replace('backend import numpy as tf',
                                 'backend import jax as tf')

tensorflow_probability/substrates/BUILD

@@ -46,12 +46,7 @@
 
 FLAGS = flags.FLAGS
 
-DO_NOT_GENERATE_DOCS_FOR = [
-    tfp.experimental.substrates.jax.tf2jax,
-    tfp.experimental.substrates.jax.experimental,
-    tfp.experimental.substrates.numpy.tf2numpy,
-    tfp.experimental.substrates.numpy.experimental,
-]
+DO_NOT_GENERATE_DOCS_FOR = []
 
 
 def internal_filter(path, parent, children):