Add tfp.experimental.math.manual_special_functions.

These trade speed for precision. They are sometimes useful when running large MCMC problems on a TPU.

PiperOrigin-RevId: 390191832

Author: SiegeLordEx

tensorflow_probability/python/build_defs.bzl

@@ -28,7 +28,11 @@ RUNFILES_ROOT = "tensorflow_probability/"
 
 def _substrate_src(src, substrate):
     """Rewrite a single src filename for the given substrate."""
-    return "_{}/_generated_{}".format(substrate, src)
+
+    # When src's are sourced from a different package we cut away the package
+    # name.
+    parts = src.split(":")
+    return "_{}/_generated_{}".format(substrate, parts[-1])
 
 def _substrate_srcs(srcs, substrate):
     """Rewrite src filenames for the given substrate."""
@@ -271,6 +275,7 @@ def multi_substrate_py_test(
         srcs = [],
         main = None,
         deps = [],
+        jax_extra_deps = [],
         tags = [],
         numpy_tags = [],
         jax_tags = [],
@@ -279,7 +284,8 @@ def multi_substrate_py_test(
         srcs_version = "PY3",
         python_version = "PY3",
         timeout = None,
-        shard_count = None):
+        shard_count = None,
+        args = []):
     """A TFP `py_test` for each of TF, NumPy, and JAX.
 
     Args:
@@ -296,6 +302,7 @@ def multi_substrate_py_test(
             use-case of the `main` argument is a secondary, i.e. GPU, test.
         deps: As with `py_test`. The list is rewritten to depend on
             substrate-specific libraries for substrate variants.
+        jax_extra_deps: Extra dependencies for the JAX substrate.
         tags: Tags global to this test target. NumPy also gets a `'tfp_numpy'`
             tag, and JAX gets a `'tfp_jax'` tag. A `f'_{name}'` tag is used
             to produce the `test_suite`.
@@ -308,6 +315,7 @@ def multi_substrate_py_test(
         python_version: As with `py_test`.
         timeout: As with `py_test`.
         shard_count: As with `py_test`.
+        args: As with `py_test`.
     """
 
     tags = [t for t in tags]
@@ -325,6 +333,7 @@ def multi_substrate_py_test(
         python_version = python_version,
         timeout = timeout,
         shard_count = shard_count,
+        args = args,
     )
     test_targets.append(":{}.tf".format(name))
 
@@ -349,6 +358,7 @@ def multi_substrate_py_test(
             python_version = "PY3",
             timeout = timeout,
             shard_count = shard_count,
+            args = args,
         )
         test_targets.append(":{}.numpy".format(name))
 
@@ -362,7 +372,7 @@ def multi_substrate_py_test(
                 cmd = "$(location {}) $(SRCS) --numpy_to_jax > $@".format(REWRITER_TARGET),
                 exec_tools = [REWRITER_TARGET],
             )
-        jax_deps = _substrate_deps(deps, "jax")
+        jax_deps = _substrate_deps(deps, "jax") + jax_extra_deps
         # [internal] Add JAX build dep
         native.py_test(
             name = "{}.jax".format(name),
@@ -375,6 +385,7 @@ def multi_substrate_py_test(
             python_version = "PY3",
             timeout = timeout,
             shard_count = shard_count,
+            args = args,
         )
         test_targets.append(":{}.jax".format(name))
 

tensorflow_probability/python/experimental/BUILD

@@ -57,6 +57,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/experimental/distributions",
         "//tensorflow_probability/python/experimental/linalg",
         "//tensorflow_probability/python/experimental/marginalize",
+        "//tensorflow_probability/python/experimental/math",
         "//tensorflow_probability/python/experimental/mcmc",
         "//tensorflow_probability/python/experimental/nn",
         "//tensorflow_probability/python/experimental/parallel_filter",

tensorflow_probability/python/experimental/__init__.py

@@ -37,6 +37,7 @@
 from tensorflow_probability.python.experimental import distributions
 from tensorflow_probability.python.experimental import linalg
 from tensorflow_probability.python.experimental import marginalize
+from tensorflow_probability.python.experimental import math
 from tensorflow_probability.python.experimental import mcmc
 from tensorflow_probability.python.experimental import nn
 from tensorflow_probability.python.experimental import parallel_filter
@@ -62,6 +63,7 @@
     'distributions',
     'linalg',
     'marginalize',
+    'math',
     'mcmc',
     'nn',
     'parallel_filter',

tensorflow_probability/python/experimental/math/BUILD

@@ -0,0 +1,72 @@
+# Copyright 2021 The TensorFlow Probability Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# Description:
+#   Experimental math.
+
+load(
+    "//tensorflow_probability/python:build_defs.bzl",
+    "multi_substrate_py_library",
+    "multi_substrate_py_test",
+)
+
+package(
+    default_visibility = [
+        "//tensorflow_probability:__subpackages__",
+    ],
+)
+
+licenses(["notice"])
+
+multi_substrate_py_library(
+    name = "math",
+    srcs = ["__init__.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":manual_special_functions",
+    ],
+)
+
+multi_substrate_py_library(
+    name = "manual_special_functions",
+    srcs = ["manual_special_functions.py"],
+    srcs_version = "PY3",
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/internal:custom_gradient",
+        "//tensorflow_probability/python/internal:dtype_util",
+    ],
+)
+
+multi_substrate_py_test(
+    name = "manual_special_functions_test",
+    srcs = ["manual_special_functions_test.py"],
+    srcs_version = "PY3",
+    tags = [
+        "tf1-broken",
+    ],
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
+exports_files(
+    glob(["**/*.py"]),
+    visibility = ["//tensorflow_probability:__subpackages__"],
+)

tensorflow_probability/python/experimental/math/__init__.py

@@ -0,0 +1,33 @@
+# Copyright 2021 The TensorFlow Probability Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Experimental math."""
+
+from tensorflow_probability.python.experimental.math.manual_special_functions import exp_pade_4_4
+from tensorflow_probability.python.experimental.math.manual_special_functions import expm1_pade_4_4
+from tensorflow_probability.python.experimental.math.manual_special_functions import log1p_pade_4_4
+from tensorflow_probability.python.experimental.math.manual_special_functions import log_pade_4_4_newton
+from tensorflow_probability.python.experimental.math.manual_special_functions import patch_manual_special_functions
+from tensorflow_probability.python.experimental.math.manual_special_functions import reduce_logsumexp
+from tensorflow_probability.python.experimental.math.manual_special_functions import softplus
+
+__all__ = [
+    'exp_pade_4_4',
+    'expm1_pade_4_4',
+    'log1p_pade_4_4',
+    'log_pade_4_4_newton',
+    'patch_manual_special_functions',
+    'reduce_logsumexp',
+    'softplus',
+]