Remove jax and numpy substrates from the main tensorflow_probability BUILD

target.

PiperOrigin-RevId: 613574866

Author: ThomasColthurst

tensorflow_probability/BUILD

@@ -34,6 +34,7 @@ licenses(["notice"])
 
 exports_files(["LICENSE"])
 
+# The public target for the TensorFlow version of TFP.
 py_library(
     name = "tensorflow_probability",
     srcs = ["__init__.py"],
@@ -42,7 +43,6 @@ py_library(
 #         "//tensorflow_probability/google/staging",  # DisableOnExport
 #         "//tensorflow_probability/google/tfp_google",  # DisableOnExport
         "//tensorflow_probability/python",
-        "//tensorflow_probability/substrates",
     ],
 )
 

tensorflow_probability/__init__.py

@@ -17,7 +17,6 @@
 # Contributors to the `python/` dir should not alter this file; instead update
 # `python/__init__.py` as necessary.
 
-from tensorflow_probability import substrates
 # from tensorflow_probability.google import staging  # DisableOnExport
 # from tensorflow_probability.google import tfp_google  # DisableOnExport
 from tensorflow_probability.python import *  # pylint: disable=wildcard-import

tensorflow_probability/substrates/meta/rewrite.py

@@ -36,8 +36,7 @@
         'from tensorflow_probability.python.internal.backend.numpy.compat '
         'import v2',
     'import tensorflow_probability as tfp':
-        'import tensorflow_probability as tfp; '
-        'tfp = tfp.substrates.numpy',
+        'from tensorflow_probability.substrates import numpy as tfp',
     ('from tensorflow.python.framework '
      'import composite_tensor'):
         ('from tensorflow_probability.python.internal.backend.numpy '