From f9ab36e4c5de43d2e9e97329ef44b03a9a282088 Mon Sep 17 00:00:00 2001 From: thomaswc Date: Fri, 26 Apr 2024 10:54:53 -0700 Subject: [PATCH] The tfp jax substrate uses jax and thus should depend on jax. PiperOrigin-RevId: 628451939 --- tensorflow_probability/substrates/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow_probability/substrates/BUILD b/tensorflow_probability/substrates/BUILD index 5a186f4cfa..5e45d3fcfb 100644 --- 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