The tfp jax substrate uses jax and thus should depend on jax.

PiperOrigin-RevId: 628451939
tensorflow · Apr 26, 2024 · f9ab36e · f9ab36e
1 parent fc1d585
commit f9ab36e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tensorflow_probability/substrates/BUILD b/tensorflow_probability/substrates/BUILD
@@ -48,6 +48,7 @@ py_library(
     srcs = ["jax/__init__.py"],
     tags = ["alt_dep=//tensorflow_probability:jax"],
     deps = [
+        # jax dep,
 #         "//tensorflow_probability/google:google.jax",  # DisableOnExport
 #         "//tensorflow_probability/google/autosts:autosts.jax",  # DisableOnExport
 #         "//tensorflow_probability/google/staging:staging.jax",  # DisableOnExport